if objective_at_initial_simplex is not None and initial_simplex is None:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if objective_at_initial_vertex is not None and initial_vertex is None:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if initial_simplex is not None:<EOL><INDENT>if initial_vertex is not None:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if step_sizes is not None:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>return _prepare_args_with_initial_simplex(objective_function,<EOL>initial_simplex,<EOL>objective_at_initial_simplex,<EOL>batch_evaluate_objective)<EOL><DEDENT>if initial_vertex is None:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if step_sizes is None:<EOL><INDENT>step_sizes = _default_step_sizes(initial_vertex)<EOL><DEDENT>return _prepare_args_with_initial_vertex(objective_function,<EOL>initial_vertex,<EOL>step_sizes,<EOL>objective_at_initial_vertex,<EOL>batch_evaluate_objective)<EOL>

Computes the initial simplex and the objective values at the simplex.

    Args:
      objective_function:  A Python callable that accepts a point as a
        real `Tensor` and returns a `Tensor` of real dtype containing
        the value of the function at that point. The function
        to be evaluated at the simplex. If `batch_evaluate_objective` is `True`,
        the callable may be evaluated on a `Tensor` of shape `[n+1] + s `
        where `n` is the dimension of the problem and `s` is the shape of a
        single point in the domain (so `n` is the size of a `Tensor`
        representing a single point).
        In this case, the expected return value is a `Tensor` of shape `[n+1]`.
      initial_simplex: None or `Tensor` of real dtype. The initial simplex to
        start the search. If supplied, should be a `Tensor` of shape `[n+1] + s`
        where `n` is the dimension of the problem and `s` is the shape of a
        single point in the domain. Each row (i.e. the `Tensor` with a given
        value of the first index) is interpreted as a vertex of a simplex and
        hence the rows must be affinely independent. If not supplied, an axes
        aligned simplex is constructed using the `initial_vertex` and
        `step_sizes`. Only one and at least one of `initial_simplex` and
        `initial_vertex` must be supplied.
      initial_vertex: None or `Tensor` of real dtype and any shape that can
        be consumed by the `objective_function`. A single point in the domain that
        will be used to construct an axes aligned initial simplex.
      step_sizes: None or `Tensor` of real dtype and shape broadcasting
        compatible with `initial_vertex`. Supplies the simplex scale along each
        axes. Only used if `initial_simplex` is not supplied. See the docstring
        of `minimize` for more details.
      objective_at_initial_simplex: None or rank `1` `Tensor` of real dtype.
        The value of the objective function at the initial simplex.
        May be supplied only if `initial_simplex` is
        supplied. If not supplied, it will be computed.
      objective_at_initial_vertex: None or scalar `Tensor` of real dtype. The
        value of the objective function at the initial vertex. May be supplied
        only if the `initial_vertex` is also supplied.
      batch_evaluate_objective: Python `bool`. If True, the objective function
        will be evaluated on all the vertices of the simplex packed into a
        single tensor. If False, the objective will be mapped across each
        vertex separately.

    Returns:
      prepared_args: A tuple containing the following elements:
        dimension: Scalar `Tensor` of `int32` dtype. The dimension of the problem
          as inferred from the supplied arguments.
        num_vertices: Scalar `Tensor` of `int32` dtype. The number of vertices
          in the simplex.
        simplex: A `Tensor` of same dtype as `initial_simplex`
          (or `initial_vertex`). The first component of the shape of the
          `Tensor` is `num_vertices` and each element represents a vertex of
          the simplex.
        objective_at_simplex: A `Tensor` of same dtype as the dtype of the
          return value of objective_function. The shape is a vector of size
          `num_vertices`. The objective function evaluated at the simplex.
        num_evaluations: An `int32` scalar `Tensor`. The number of points on
          which the objective function was evaluated.

    Raises:
      ValueError: If any of the following conditions hold
        1. If none or more than one of `initial_simplex` and `initial_vertex` are
          supplied.
        2. If `initial_simplex` and `step_sizes` are both specified.


<EOL>small_sizes = tf.ones_like(reference_vertex) * <NUM_LIT><EOL>large_sizes = reference_vertex * <NUM_LIT><EOL>return tf.where(tf.abs(reference_vertex) < _EPSILON,<EOL>small_sizes,<EOL>large_sizes)<EOL>

Chooses default step sizes according to [Gao and Han(2010)][3].


initial_simplex = tf.convert_to_tensor(value=initial_simplex)<EOL>num_vertices = tf.shape(input=initial_simplex)[<NUM_LIT:0>]<EOL>dim = num_vertices - <NUM_LIT:1><EOL>num_evaluations = <NUM_LIT:0><EOL>if objective_at_initial_simplex is None:<EOL><INDENT>objective_at_initial_simplex, n_evals = _evaluate_objective_multiple(<EOL>objective_function, initial_simplex, batch_evaluate_objective)<EOL>num_evaluations += n_evals<EOL><DEDENT>objective_at_initial_simplex = tf.convert_to_tensor(<EOL>value=objective_at_initial_simplex)<EOL>return (dim,<EOL>num_vertices,<EOL>initial_simplex,<EOL>objective_at_initial_simplex,<EOL>num_evaluations)<EOL>

Evaluates the objective function at the specified initial simplex.


dim = tf.size(input=initial_vertex)<EOL>num_vertices = dim + <NUM_LIT:1><EOL>unit_vectors_along_axes = tf.reshape(<EOL>tf.eye(dim, dim, dtype=initial_vertex.dtype.base_dtype),<EOL>tf.concat([[dim], tf.shape(input=initial_vertex)], axis=<NUM_LIT:0>))<EOL>simplex_face = initial_vertex + step_sizes * unit_vectors_along_axes<EOL>simplex = tf.concat([tf.expand_dims(initial_vertex, axis=<NUM_LIT:0>),<EOL>simplex_face], axis=<NUM_LIT:0>)<EOL>num_evaluations = <NUM_LIT:0><EOL>if objective_at_initial_vertex is None:<EOL><INDENT>objective_at_initial_vertex = objective_function(initial_vertex)<EOL>num_evaluations += <NUM_LIT:1><EOL><DEDENT>objective_at_simplex_face, num_evals = _evaluate_objective_multiple(<EOL>objective_function, simplex_face, batch_evaluate_objective)<EOL>num_evaluations += num_evals<EOL>objective_at_simplex = tf.concat(<EOL>[<EOL>tf.expand_dims(objective_at_initial_vertex, axis=<NUM_LIT:0>),<EOL>objective_at_simplex_face<EOL>], axis=<NUM_LIT:0>)<EOL>return (dim,<EOL>num_vertices,<EOL>simplex,<EOL>objective_at_simplex,<EOL>num_evaluations)<EOL>

Constructs a standard axes aligned simplex.


dim = tf.cast(dim, dtype=dtype)<EOL>reflection = <NUM_LIT:1.> if reflection is None else reflection<EOL>expansion = (<NUM_LIT:1.> + <NUM_LIT> / dim) if expansion is None else expansion<EOL>contraction = (<NUM_LIT> - <NUM_LIT:1.> / (<NUM_LIT:2> * dim)) if contraction is None else contraction<EOL>shrinkage = (<NUM_LIT:1.> - <NUM_LIT:1.> / dim) if shrinkage is None else shrinkage<EOL>return reflection, expansion, contraction, shrinkage<EOL>

Applies the [Gao and Han][3] presciption to the unspecified parameters.


n_points = tf.shape(input=arg_batch)[<NUM_LIT:0>]<EOL>if batch_evaluate_objective:<EOL><INDENT>return objective_function(arg_batch), n_points<EOL><DEDENT>return tf.map_fn(objective_function, arg_batch), n_points<EOL>

Evaluates the objective function on a batch of points.

    If `batch_evaluate_objective` is True, returns
    `objective function(arg_batch)` else it maps the `objective_function`
    across the `arg_batch`.

    Args:
      objective_function: A Python callable that accepts a single `Tensor` of
        rank 'R > 1' and any shape 's' and returns a scalar `Tensor` of real dtype
        containing the value of the function at that point. If
        `batch a `Tensor` of shape `[batch_size] + s ` where `batch_size` is the
        size of the batch of args. In this case, the expected return value is a
        `Tensor` of shape `[batch_size]`.
      arg_batch: A `Tensor` of real dtype. The batch of arguments at which to
        evaluate the `objective_function`. If `batch_evaluate_objective` is False,
        `arg_batch` will be unpacked along the zeroth axis and the
        `objective_function` will be applied to each element.
      batch_evaluate_objective: `bool`. Whether the `objective_function` can
        evaluate a batch of arguments at once.

    Returns:
      A tuple containing:
        objective_values: A `Tensor` of real dtype and shape `[batch_size]`.
          The value of the objective function evaluated at the supplied
          `arg_batch`.
        num_evaluations: An `int32` scalar `Tensor`containing the number of
          points on which the objective function was evaluated (i.e `batch_size`).


return (tf.reduce_any(input_tensor=converged) |<EOL>tf.reduce_all(input_tensor=failed))<EOL>

Condition to stop when any batch member converges, or all have failed.


return tf.reduce_all(input_tensor=converged | failed)<EOL>

Condition to stop when all batch members have converged or failed.


if control_inputs:<EOL><INDENT>with tf.control_dependencies(control_inputs):<EOL><INDENT>f0, df0 = value_and_gradients_function(initial_position)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>f0, df0 = value_and_gradients_function(initial_position)<EOL><DEDENT>converged = norm(df0, dims=<NUM_LIT:1>) < grad_tolerance<EOL>return dict(<EOL>converged=converged,<EOL>failed=tf.zeros_like(converged),  <EOL>num_iterations=tf.convert_to_tensor(value=<NUM_LIT:0>),<EOL>num_objective_evaluations=tf.convert_to_tensor(value=<NUM_LIT:1>),<EOL>position=initial_position,<EOL>objective_value=f0,<EOL>objective_gradient=df0)<EOL>

Returns a dictionary to populate the initial state of the search procedure.

    Performs an initial convergence check and the first evaluation of the
    objective function.

    Args:
      value_and_gradients_function: A Python callable that accepts a tensor and
        returns a tuple of two tensors: the objective function value and its
        derivative.
      initial_position: The starting point of the search procedure.
      grad_tolerance: The gradient tolerance for the procedure.
      control_inputs: Optional ops used to assert the validity of inputs, these
        are added as control dependencies to execute before the objective
        function is evaluated for the first time.

    Returns:
      An dictionary with values for the following keys:
        converged: True if the convergence check finds that the initial position
          is already an argmin of the objective function.
        failed: Initialized to False.
        num_objective_evaluations: Initialized to 1.
        position: Initialized to the initial position.
        objective_value: Initialized to the value of the objective function at
          the initial position.
        objective_gradient: Initialized to the gradient of the objective
          function at the initial position.


def line_search_step(state, value_and_gradients_function, search_direction,<EOL>grad_tolerance, f_relative_tolerance, x_tolerance,<EOL>stopping_condition):

line_search_value_grad_func = _restrict_along_direction(<EOL>value_and_gradients_function, state.position, search_direction)<EOL>derivative_at_start_pt = tf.reduce_sum(<EOL>input_tensor=state.objective_gradient * search_direction, axis=-<NUM_LIT:1>)<EOL>val_0 = ValueAndGradient(x=_broadcast(<NUM_LIT:0>, state.position),<EOL>f=state.objective_value,<EOL>df=derivative_at_start_pt,<EOL>full_gradient=state.objective_gradient)<EOL>inactive = state.failed | state.converged<EOL>ls_result = linesearch.hager_zhang(<EOL>line_search_value_grad_func,<EOL>initial_step_size=_broadcast(<NUM_LIT:1>, state.position),<EOL>value_at_zero=val_0,<EOL>converged=inactive)  <EOL>state_after_ls = update_fields(<EOL>state,<EOL>failed=state.failed | ~ls_result.converged,<EOL>num_iterations=state.num_iterations + <NUM_LIT:1>,<EOL>num_objective_evaluations=(<EOL>state.num_objective_evaluations + ls_result.func_evals))<EOL>def _do_update_position():<EOL><INDENT>position_delta = tf.where(<EOL>inactive,<EOL>tf.zeros_like(search_direction),<EOL>search_direction * tf.expand_dims(ls_result.left.x, axis=-<NUM_LIT:1>))<EOL>return _update_position(<EOL>state_after_ls,<EOL>position_delta,<EOL>ls_result.left.f,<EOL>ls_result.left.full_gradient,<EOL>grad_tolerance, f_relative_tolerance, x_tolerance)<EOL><DEDENT>return prefer_static.cond(<EOL>stopping_condition(state.converged, state.failed),<EOL>true_fn=lambda: state_after_ls,<EOL>false_fn=_do_update_position)<EOL>

Performs the line search step of the BFGS search procedure.

    Uses hager_zhang line search procedure to compute a suitable step size
    to advance the current `state.position` along the given `search_direction`.
    Also, if the line search is successful, updates the `state.position` by
    taking the corresponding step.

    Args:
      state: A namedtuple instance holding values for the current state of the
        search procedure. The state must include the fields: `position`,
        `objective_value`, `objective_gradient`, `num_iterations`,
        `num_objective_evaluations`, `converged` and `failed`.
      value_and_gradients_function: A Python callable that accepts a point as a
        real `Tensor` of shape `[..., n]` and returns a tuple of two tensors of
        the same dtype: the objective function value, a real `Tensor` of shape
        `[...]`, and its derivative, another real `Tensor` of shape `[..., n]`.
      search_direction: A real `Tensor` of shape `[..., n]`. The direction along
        which to perform line search.
      grad_tolerance: Scalar `Tensor` of real dtype. Specifies the gradient
        tolerance for the procedure.
      f_relative_tolerance: Scalar `Tensor` of real dtype. Specifies the
        tolerance for the relative change in the objective value.
      x_tolerance: Scalar `Tensor` of real dtype. Specifies the tolerance for the
        change in the position.
      stopping_condition: A Python function that takes as input two Boolean
        tensors of shape `[...]`, and returns a Boolean scalar tensor. The input
        tensors are `converged` and `failed`, indicating the current status of
        each respective batch member; the return value states whether the
        algorithm should stop.
    Returns:
      A copy of the input state with the following fields updated:
        converged: a Boolean `Tensor` of shape `[...]` indicating whether the
          convergence criteria has been met.
        failed: a Boolean `Tensor` of shape `[...]` indicating whether the line
          search procedure failed to converge, or if either the updated gradient
          or objective function are no longer finite.
        num_iterations: Increased by 1.
        num_objective_evaluations: Increased by the number of times that the
          objective function got evaluated.
        position, objective_value, objective_gradient: If line search succeeded,
          updated by computing the new position and evaluating the objective
          function at that position.


return state._replace(**kwargs)<EOL>

Copies the argument and overrides some of its fields.

    Args:
      state: A `collections.namedtuple` instance.
      **kwargs: Other named arguments represent fields in the tuple to override
        with new values.

    Returns:
      A namedtuple, of the same class as the input argument, with the updated
      fields.

    Raises:
      ValueError if the supplied kwargs contain fields not present in the
      input argument.


def _restricted_func(t):<EOL><INDENT>t = _broadcast(t, position)<EOL>pt = position + tf.expand_dims(t, axis=-<NUM_LIT:1>) * direction<EOL>objective_value, gradient = value_and_gradients_function(pt)<EOL>return ValueAndGradient(<EOL>x=t,<EOL>f=objective_value,<EOL>df=tf.reduce_sum(input_tensor=gradient * direction, axis=-<NUM_LIT:1>),<EOL>full_gradient=gradient)<EOL><DEDENT>return _restricted_func<EOL>

Restricts a function in n-dimensions to a given direction.

    Suppose f: R^n -> R. Then given a point x0 and a vector p0 in R^n, the
    restriction of the function along that direction is defined by:

    ```None
    g(t) = f(x0 + t * p0)
    ```

    This function performs this restriction on the given function. In addition, it
    also computes the gradient of the restricted function along the restriction
    direction. This is equivalent to computing `dg/dt` in the definition above.

    Args:
      value_and_gradients_function: Callable accepting a single real `Tensor`
        argument of shape `[..., n]` and returning a tuple of a real `Tensor` of
        shape `[...]` and a real `Tensor` of shape `[..., n]`. The multivariate
        function whose restriction is to be computed. The output values of the
        callable are the function value and the gradients at the input argument.
      position: `Tensor` of real dtype and shape consumable by
        `value_and_gradients_function`. Corresponds to `x0` in the definition
        above.
      direction: `Tensor` of the same dtype and shape as `position`. The direction
        along which to restrict the function. Note that the direction need not
        be a unit vector.

    Returns:
      restricted_value_and_gradients_func: A callable accepting a tensor of shape
        broadcastable to `[...]` and same dtype as `position` and returning a
        namedtuple of `Tensors`. The input tensor is the parameter along the
        direction labelled `t` above. The return value contains fields:
          x: A real `Tensor` of shape `[...]`. The input value `t` where the line
            function was evaluated, after any necessary broadcasting.
          f: A real `Tensor` of shape `[...]` containing the value of the
            function at the point `position + t * direction`.
          df: A real `Tensor` of shape `[...]` containing the derivative at
            `position + t * direction`.
          full_gradient: A real `Tensor` of shape `[..., n]`, the full gradient
            of the original `value_and_gradients_function`.


failed = state.failed | ~tf.math.is_finite(next_objective) | ~tf.reduce_all(<EOL>input_tensor=tf.math.is_finite(next_gradient), axis=-<NUM_LIT:1>)<EOL>next_position = state.position + position_delta<EOL>converged = ~failed & _check_convergence(state.position,<EOL>next_position,<EOL>state.objective_value,<EOL>next_objective,<EOL>next_gradient,<EOL>grad_tolerance,<EOL>f_relative_tolerance,<EOL>x_tolerance)<EOL>return update_fields(<EOL>state,<EOL>converged=state.converged | converged,<EOL>failed=failed,<EOL>position=next_position,<EOL>objective_value=next_objective,<EOL>objective_gradient=next_gradient)<EOL>

Updates the state advancing its position by a given position_delta.


if dims == <NUM_LIT:0>:<EOL><INDENT>return tf.math.abs(value)<EOL><DEDENT>elif dims == <NUM_LIT:1>:<EOL><INDENT>axis = -<NUM_LIT:1><EOL><DEDENT>elif dims == <NUM_LIT:2>:<EOL><INDENT>axis = [-<NUM_LIT:1>, -<NUM_LIT:2>]<EOL><DEDENT>else:<EOL><INDENT>ValueError(dims)<EOL><DEDENT>if order is None:<EOL><INDENT>order = np.inf<EOL><DEDENT>return tf.norm(tensor=value, axis=axis, ord=order)<EOL>

Compute the norm of the given (possibly batched) value.

    Args:
      value: A `Tensor` of real dtype.
      dims: An Python integer with the number of non-batching dimensions in the
        value, i.e. `dims=0` (scalars), `dims=1` (vectors), `dims=2` (matrices).
      order: Order of the norm, defaults to `np.inf`.


grad_converged = norm(next_gradient, dims=<NUM_LIT:1>) <= grad_tolerance<EOL>x_converged = norm(next_position - current_position, dims=<NUM_LIT:1>) <= x_tolerance<EOL>f_converged = (norm(next_objective - current_objective, dims=<NUM_LIT:0>) <=<EOL>f_relative_tolerance * current_objective)<EOL>return grad_converged | x_converged | f_converged<EOL>

Checks if the algorithm satisfies the convergence criteria.


return tf.broadcast_to(<EOL>tf.convert_to_tensor(value=value, dtype=target.dtype),<EOL>distribution_util.prefer_static_shape(target)[:-<NUM_LIT:1>])<EOL>

Broadcast a value to match the batching dimensions of a target.

    If necessary the value is converted into a tensor. Both value and target
    should be of the same dtype.

    Args:
      value: A value to broadcast.
      target: A `Tensor` of shape [b1, ..., bn, d].

    Returns:
      A `Tensor` of shape [b1, ..., bn] and same dtype as the target.


if isinstance(sparse_or_dense_matrix,<EOL>(tf.SparseTensor, tf.compat.v1.SparseTensorValue)):<EOL><INDENT>num_rows = _get_shape(sparse_or_dense_matrix)[-<NUM_LIT:2>]<EOL>batch_shape = _get_shape(sparse_or_dense_matrix)[:-<NUM_LIT:2>]<EOL>slice_start = tf.concat([tf.zeros_like(batch_shape), [<NUM_LIT:0>, col_index]],<EOL>axis=<NUM_LIT:0>)<EOL>slice_size = tf.concat([batch_shape, [num_rows, <NUM_LIT:1>]], axis=<NUM_LIT:0>)<EOL>sparse_slice = tf.sparse.slice(sparse_or_dense_matrix,<EOL>tf.cast(slice_start, tf.int64),<EOL>tf.cast(slice_size, tf.int64))<EOL>output_shape = tf.concat([batch_shape, [num_rows]], axis=<NUM_LIT:0>)<EOL>return tf.reshape(tf.sparse.to_dense(sparse_slice), output_shape)<EOL><DEDENT>else:<EOL><INDENT>return tf.gather(sparse_or_dense_matrix, col_index, axis=-<NUM_LIT:1>)<EOL><DEDENT>

Returns a (dense) column of a Tensor or SparseTensor.

    Args:
      sparse_or_dense_matrix: matrix-shaped, `float` `Tensor` or `SparseTensor`.
      col_index: scalar, `int` `Tensor` representing the index of the desired
        column.

    Returns:
      column: vector-shaped, `float` `Tensor` with the same dtype as
        `sparse_or_dense_matrix`, representing the `col_index`th column of
        `sparse_or_dense_matrix`.


graph_deps = [<EOL>gradient_unregularized_loss,<EOL>hessian_unregularized_loss_outer,<EOL>hessian_unregularized_loss_middle,<EOL>x_start,<EOL>l1_regularizer,<EOL>l2_regularizer,<EOL>maximum_full_sweeps,<EOL>tolerance,<EOL>learning_rate,<EOL>]<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>', graph_deps):<EOL><INDENT>x_shape = _get_shape(x_start)<EOL>batch_shape = x_shape[:-<NUM_LIT:1>]<EOL>dims = x_shape[-<NUM_LIT:1>]<EOL>def _hessian_diag_elt_with_l2(coord):  <EOL><INDENT>inner_square = tf.reduce_sum(<EOL>input_tensor=_sparse_or_dense_matmul_onehot(<EOL>hessian_unregularized_loss_outer, coord)**<NUM_LIT:2>,<EOL>axis=-<NUM_LIT:1>)<EOL>unregularized_component = (<EOL>hessian_unregularized_loss_middle[..., coord] * inner_square)<EOL>l2_component = _mul_or_none(<NUM_LIT>, l2_regularizer)<EOL>return _add_ignoring_nones(unregularized_component, l2_component)<EOL><DEDENT>grad_loss_with_l2 = _add_ignoring_nones(<EOL>gradient_unregularized_loss, _mul_or_none(<NUM_LIT>, l2_regularizer, x_start))<EOL>x_update_diff_norm_sq_convergence_threshold = (<EOL>tolerance * (<NUM_LIT:1.> + tf.norm(tensor=x_start, ord=<NUM_LIT:2>, axis=-<NUM_LIT:1>))**<NUM_LIT:2>)<EOL>update_shape = tf.concat([[dims], batch_shape], axis=-<NUM_LIT:1>)<EOL>def _loop_cond(iter_, x_update_diff_norm_sq, x_update,<EOL>hess_matmul_x_update):<EOL><INDENT>del x_update<EOL>del hess_matmul_x_update<EOL>sweep_complete = (iter_ > <NUM_LIT:0>) & tf.equal(iter_ % dims, <NUM_LIT:0>)<EOL>small_delta = (<EOL>x_update_diff_norm_sq < x_update_diff_norm_sq_convergence_threshold)<EOL>converged = sweep_complete & small_delta<EOL>allowed_more_iterations = iter_ < maximum_full_sweeps * dims<EOL>return allowed_more_iterations & tf.reduce_any(input_tensor=~converged)<EOL><DEDENT>def _loop_body(  <EOL>iter_, x_update_diff_norm_sq, x_update, hess_matmul_x_update):<EOL><INDENT>coord = iter_ % dims<EOL>x_update_diff_norm_sq = tf.where(<EOL>tf.equal(coord, <NUM_LIT:0>), tf.zeros_like(x_update_diff_norm_sq),<EOL>x_update_diff_norm_sq)<EOL>w_old = x_start[..., coord] + x_update[coord, ...]<EOL>second_deriv = _hessian_diag_elt_with_l2(coord)<EOL>newton_step = -_mul_ignoring_nones(  <EOL>learning_rate, grad_loss_with_l2[..., coord] +<EOL>hess_matmul_x_update[coord, ...]) / second_deriv<EOL>delta = (<EOL>soft_threshold(<EOL>w_old + newton_step,<EOL>_mul_ignoring_nones(learning_rate, l1_regularizer) / second_deriv)<EOL>- w_old)<EOL>def _do_update(x_update_diff_norm_sq, x_update, hess_matmul_x_update):  <EOL><INDENT>hessian_column_with_l2 = sparse_or_dense_matvecmul(<EOL>hessian_unregularized_loss_outer,<EOL>hessian_unregularized_loss_middle * _sparse_or_dense_matmul_onehot(<EOL>hessian_unregularized_loss_outer, coord),<EOL>adjoint_a=True)<EOL>if l2_regularizer is not None:<EOL><INDENT>hessian_column_with_l2 += _one_hot_like(<EOL>hessian_column_with_l2, coord, on_value=<NUM_LIT> * l2_regularizer)<EOL><DEDENT>n = tf.rank(hessian_column_with_l2)<EOL>perm = tf.roll(tf.range(n), shift=<NUM_LIT:1>, axis=<NUM_LIT:0>)<EOL>hessian_column_with_l2 = tf.transpose(<EOL>a=hessian_column_with_l2, perm=perm)<EOL>x_update = tf.tensor_scatter_nd_add(x_update, [[coord]], [delta])<EOL>with tf.control_dependencies([x_update]):<EOL><INDENT>x_update_diff_norm_sq_ = x_update_diff_norm_sq + delta**<NUM_LIT:2><EOL>hess_matmul_x_update_ = (<EOL>hess_matmul_x_update + delta * hessian_column_with_l2)<EOL>x_update_diff_norm_sq_.set_shape(<EOL>x_update_diff_norm_sq_.shape.merge_with(<EOL>x_update_diff_norm_sq.shape))<EOL>hess_matmul_x_update_.set_shape(<EOL>hess_matmul_x_update_.shape.merge_with(<EOL>hess_matmul_x_update.shape))<EOL>return [x_update_diff_norm_sq_, x_update, hess_matmul_x_update_]<EOL><DEDENT><DEDENT>inputs_to_update = [x_update_diff_norm_sq, x_update, hess_matmul_x_update]<EOL>return [iter_ + <NUM_LIT:1>] + prefer_static.cond(<EOL>tf.reduce_all(input_tensor=tf.equal(delta, <NUM_LIT:0.>)),<EOL>lambda: inputs_to_update,<EOL>lambda: _do_update(*inputs_to_update))<EOL><DEDENT>base_dtype = x_start.dtype.base_dtype<EOL>iter_, x_update_diff_norm_sq, x_update, _ = tf.while_loop(<EOL>cond=_loop_cond,<EOL>body=_loop_body,<EOL>loop_vars=[<EOL>tf.zeros([], dtype=np.int32, name='<STR_LIT>'),<EOL>tf.zeros(<EOL>batch_shape, dtype=base_dtype, name='<STR_LIT>'),<EOL>tf.zeros(update_shape, dtype=base_dtype, name='<STR_LIT>'),<EOL>tf.zeros(<EOL>update_shape, dtype=base_dtype, name='<STR_LIT>'),<EOL>])<EOL>n = tf.rank(x_update)<EOL>perm = tf.roll(tf.range(n), shift=-<NUM_LIT:1>, axis=<NUM_LIT:0>)<EOL>x_update = tf.transpose(a=x_update, perm=perm)<EOL>converged = tf.reduce_all(input_tensor=x_update_diff_norm_sq <<EOL>x_update_diff_norm_sq_convergence_threshold)<EOL>return x_start + x_update, converged, iter_ / dims<EOL><DEDENT>

One step of (the outer loop of) the minimization algorithm.

    This function returns a new value of `x`, equal to `x_start + x_update`.  The
    increment `x_update in R^n` is computed by a coordinate descent method, that
    is, by a loop in which each iteration updates exactly one coordinate of
    `x_update`.  (Some updates may leave the value of the coordinate unchanged.)

    The particular update method used is to apply an L1-based proximity operator,
    "soft threshold", whose fixed point `x_update_fix` is the desired minimum

    ```none
    x_update_fix = argmin{
        Loss(x_start + x_update')
          + l1_regularizer * ||x_start + x_update'||_1
          + l2_regularizer * ||x_start + x_update'||_2**2
        : x_update' }
    ```

    where in each iteration `x_update'` is constrained to have at most one nonzero
    coordinate.

    This update method preserves sparsity, i.e., tends to find sparse solutions if
    `x_start` is sparse.  Additionally, the choice of step size is based on
    curvature (Hessian), which significantly speeds up convergence.

    This algorithm assumes that `Loss` is convex, at least in a region surrounding
    the optimum.  (If `l2_regularizer > 0`, then only weak convexity is needed.)

    Args:
      gradient_unregularized_loss: (Batch of) `Tensor` with the same shape and
        dtype as `x_start` representing the gradient, evaluated at `x_start`, of
        the unregularized loss function (denoted `Loss` above).  (In all current
        use cases, `Loss` is the negative log likelihood.)
      hessian_unregularized_loss_outer: (Batch of) `Tensor` or `SparseTensor`
        having the same dtype as `x_start`, and shape `[N, n]` where `x_start` has
        shape `[n]`, satisfying the property
        `Transpose(hessian_unregularized_loss_outer)
        @ diag(hessian_unregularized_loss_middle)
        @ hessian_unregularized_loss_inner
        = (approximation of) Hessian matrix of Loss, evaluated at x_start`.
      hessian_unregularized_loss_middle: (Batch of) vector-shaped `Tensor` having
        the same dtype as `x_start`, and shape `[N]` where
        `hessian_unregularized_loss_outer` has shape `[N, n]`, satisfying the
        property
        `Transpose(hessian_unregularized_loss_outer)
        @ diag(hessian_unregularized_loss_middle)
        @ hessian_unregularized_loss_inner
        = (approximation of) Hessian matrix of Loss, evaluated at x_start`.
      x_start: (Batch of) vector-shaped, `float` `Tensor` representing the current
        value of the argument to the Loss function.
      tolerance: scalar, `float` `Tensor` representing the convergence threshold.
        The optimization step will terminate early, returning its current value of
        `x_start + x_update`, once the following condition is met:
        `||x_update_end - x_update_start||_2 / (1 + ||x_start||_2)
        < sqrt(tolerance)`,
        where `x_update_end` is the value of `x_update` at the end of a sweep and
        `x_update_start` is the value of `x_update` at the beginning of that
        sweep.
      l1_regularizer: scalar, `float` `Tensor` representing the weight of the L1
        regularization term (see equation above).  If L1 regularization is not
        required, then `tfp.glm.fit_one_step` is preferable.
      l2_regularizer: scalar, `float` `Tensor` representing the weight of the L2
        regularization term (see equation above).
        Default value: `None` (i.e., no L2 regularization).
      maximum_full_sweeps: Python integer specifying maximum number of sweeps to
        run.  A "sweep" consists of an iteration of coordinate descent on each
        coordinate. After this many sweeps, the algorithm will terminate even if
        convergence has not been reached.
        Default value: `1`.
      learning_rate: scalar, `float` `Tensor` representing a multiplicative factor
        used to dampen the proximal gradient descent steps.
        Default value: `None` (i.e., factor is conceptually `1`).
      name: Python string representing the name of the TensorFlow operation.
        The default name is `"minimize_one_step"`.

    Returns:
      x: (Batch of) `Tensor` having the same shape and dtype as `x_start`,
        representing the updated value of `x`, that is, `x_start + x_update`.
      is_converged: scalar, `bool` `Tensor` indicating whether convergence
        occurred across all batches within the specified number of sweeps.
      iter: scalar, `int` `Tensor` representing the actual number of coordinate
        updates made (before achieving convergence).  Since each sweep consists of
        `tf.size(x_start)` iterations, the maximum number of updates is
        `maximum_full_sweeps * tf.size(x_start)`.

    #### References

    [1]: Jerome Friedman, Trevor Hastie and Rob Tibshirani. Regularization Paths
         for Generalized Linear Models via Coordinate Descent. _Journal of
         Statistical Software_, 33(1), 2010.
         https://www.jstatsoft.org/article/view/v033i01/v33i01.pdf

    [2]: Guo-Xun Yuan, Chia-Hua Ho and Chih-Jen Lin. An Improved GLMNET for
         L1-regularized Logistic Regression. _Journal of Machine Learning
         Research_, 13, 2012.
         http://www.jmlr.org/papers/volume13/yuan12a/yuan12a.pdf


graph_deps = [<EOL>x_start,<EOL>l1_regularizer,<EOL>l2_regularizer,<EOL>maximum_iterations,<EOL>maximum_full_sweeps_per_iteration,<EOL>tolerance,<EOL>learning_rate,<EOL>],<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>', graph_deps):<EOL><INDENT>def _loop_cond(x_start, converged, iter_):<EOL><INDENT>del x_start<EOL>return tf.logical_and(iter_ < maximum_iterations,<EOL>tf.logical_not(converged))<EOL><DEDENT>def _loop_body(x_start, converged, iter_):  <EOL><INDENT>g, h_outer, h_middle = grad_and_hessian_loss_fn(x_start)<EOL>x_start, converged, _ = minimize_one_step(<EOL>gradient_unregularized_loss=g,<EOL>hessian_unregularized_loss_outer=h_outer,<EOL>hessian_unregularized_loss_middle=h_middle,<EOL>x_start=x_start,<EOL>l1_regularizer=l1_regularizer,<EOL>l2_regularizer=l2_regularizer,<EOL>maximum_full_sweeps=maximum_full_sweeps_per_iteration,<EOL>tolerance=tolerance,<EOL>learning_rate=learning_rate)<EOL>return x_start, converged, iter_ + <NUM_LIT:1><EOL><DEDENT>return tf.while_loop(<EOL>cond=_loop_cond,<EOL>body=_loop_body,<EOL>loop_vars=[<EOL>x_start,<EOL>tf.zeros([], np.bool, name='<STR_LIT>'),<EOL>tf.zeros([], np.int32, name='<STR_LIT>'),<EOL>])<EOL><DEDENT>

Minimize using Hessian-informed proximal gradient descent.

    This function solves the regularized minimization problem

    ```none
    argmin{ Loss(x)
              + l1_regularizer * ||x||_1
              + l2_regularizer * ||x||_2**2
            : x in R^n }
    ```

    where `Loss` is a convex C^2 function (typically, `Loss` is the negative log
    likelihood of a model and `x` is a vector of model coefficients).  The `Loss`
    function does not need to be supplied directly, but this optimizer does need a
    way to compute the gradient and Hessian of the Loss function at a given value
    of `x`.  The gradient and Hessian are often computationally expensive, and
    this optimizer calls them relatively few times compared with other algorithms.

    Args:
      grad_and_hessian_loss_fn: callable that takes as input a (batch of) `Tensor`
        of the same shape and dtype as `x_start` and returns the triple
        `(gradient_unregularized_loss, hessian_unregularized_loss_outer,
        hessian_unregularized_loss_middle)` as defined in the argument spec of
        `minimize_one_step`.
      x_start: (Batch of) vector-shaped, `float` `Tensor` representing the initial
        value of the argument to the `Loss` function.
      tolerance: scalar, `float` `Tensor` representing the tolerance for each
        optimization step; see the `tolerance` argument of
        `minimize_one_step`.
      l1_regularizer: scalar, `float` `Tensor` representing the weight of the L1
        regularization term (see equation above).
      l2_regularizer: scalar, `float` `Tensor` representing the weight of the L2
        regularization term (see equation above).
        Default value: `None` (i.e., no L2 regularization).
      maximum_iterations: Python integer specifying the maximum number of
        iterations of the outer loop of the optimizer.  After this many iterations
        of the outer loop, the algorithm will terminate even if the return value
        `optimal_x` has not converged.
        Default value: `1`.
      maximum_full_sweeps_per_iteration: Python integer specifying the maximum
        number of sweeps allowed in each iteration of the outer loop of the
        optimizer.  Passed as the `maximum_full_sweeps` argument to
        `minimize_one_step`.
        Default value: `1`.
      learning_rate: scalar, `float` `Tensor` representing a multiplicative factor
        used to dampen the proximal gradient descent steps.
        Default value: `None` (i.e., factor is conceptually `1`).
      name: Python string representing the name of the TensorFlow operation.
        The default name is `"minimize"`.

    Returns:
      x: `Tensor` of the same shape and dtype as `x_start`, representing the
        (batches of) computed values of `x` which minimizes `Loss(x)`.
      is_converged: scalar, `bool` `Tensor` indicating whether the minimization
        procedure converged within the specified number of iterations across all
        batches.  Here convergence means that an iteration of the inner loop
        (`minimize_one_step`) returns `True` for its `is_converged` output value.
      iter: scalar, `int` `Tensor` indicating the actual number of iterations of
        the outer loop of the optimizer completed (i.e., number of calls to
        `minimize_one_step` before achieving convergence).

    #### References

    [1]: Jerome Friedman, Trevor Hastie and Rob Tibshirani. Regularization Paths
         for Generalized Linear Models via Coordinate Descent. _Journal of
         Statistical Software_, 33(1), 2010.
         https://www.jstatsoft.org/article/view/v033i01/v33i01.pdf

    [2]: Guo-Xun Yuan, Chia-Hua Ho and Chih-Jen Lin. An Improved GLMNET for
         L1-regularized Logistic Regression. _Journal of Machine Learning
         Research_, 13, 2012.
         http://www.jmlr.org/papers/volume13/yuan12a/yuan12a.pdf


rng = np.random.RandomState(<NUM_LIT:0>)<EOL>num_points = np.prod(grid_spec.shape)<EOL>grid = np.linspace(grid_spec.min, grid_spec.max, num=num_points).astype(dtype)<EOL>grid_spacing = (grid_spec.max - grid_spec.min) / num_points<EOL>grid += <NUM_LIT:0.1> * grid_spacing * rng.randn(*grid.shape)<EOL>grid = np.sort(grid)<EOL>return np.reshape(grid, grid_spec.shape)<EOL>

Returns a uniform grid + noise, reshaped to shape argument.


<EOL>p = tf.constant(<EOL>np.array([<EOL><NUM_LIT:0.>,<EOL>np.exp(-<NUM_LIT>),<EOL>np.exp(-<NUM_LIT>),<EOL><NUM_LIT:1.> - np.exp(-<NUM_LIT>),<EOL><NUM_LIT:1.> - np.exp(-<NUM_LIT>),<EOL><NUM_LIT:1.>,<EOL>]).astype(dtype))<EOL>_, grads = value_and_gradient(special_math.ndtri, p)<EOL>self.assertAllFinite(self.evaluate(grads[<NUM_LIT:0>]))<EOL>

Verifies that ndtri has finite gradients at interesting points.


from_parts = tf.nest.flatten(from_structure)<EOL>if len(from_parts) == <NUM_LIT:1>:<EOL><INDENT>from_structure = tf.nest.map_structure(lambda _: from_parts[<NUM_LIT:0>],<EOL>to_structure)<EOL><DEDENT>return from_structure<EOL>

Broadcasts `from_structure` to `to_structure`.

    This is useful for downstream usage of `zip` or `tf.nest.map_structure`.

    If `from_structure` is a singleton, it is tiled to match the structure of
    `to_structure`. Note that the elements in `from_structure` are not copied if
    this tiling occurs.

    Args:
      to_structure: A structure.
      from_structure: A structure.

    Returns:
      new_from_structure: Same structure as `to_structure`.

    #### Example:

    ```python
    a_structure = ['a', 'b', 'c']
    b_structure = broadcast_structure(a_structure, 'd')
    # -> ['d', 'd', 'd']
    c_structure = tf.nest.map_structure(
        lambda a, b: a + b, a_structure, b_structure)
    # -> ['ad', 'bd', 'cd']
    ```


return (isinstance(args, collections.Sequence) and<EOL>not _is_namedtuple(args) and not _force_leaf(args))<EOL>

Returns `True` if `args` should be expanded as `*args`.


def _nested_convert_to_tensor(struct, dtype=None, name=None):

if dtype is not None or not tf.nest.is_nested(struct):<EOL><INDENT>return tf.convert_to_tensor(struct, dtype=dtype)<EOL><DEDENT>if _maybe_convertible_to_tensor(struct):<EOL><INDENT>try:<EOL><INDENT>return tf.convert_to_tensor(value=struct, name=name)<EOL><DEDENT>except (ValueError, TypeError):<EOL><INDENT>pass<EOL><DEDENT><DEDENT>shallow_struct = _get_shallow_structure(struct)<EOL>return nest.map_structure_up_to(<EOL>shallow_struct, lambda s: _nested_convert_to_tensor(s, name=name), struct)<EOL>

Eagerly converts struct to Tensor, recursing upon failure.


def convert_args_to_tensor(args, dtype=None, name=None):

if dtype is None:<EOL><INDENT>if expand_as_args(args) or _expand_as_kwargs(args):<EOL><INDENT>shallow_args = _get_shallow_structure(args)<EOL>return nest.map_structure_up_to(<EOL>shallow_args, lambda s: _nested_convert_to_tensor(s, name=name), args)<EOL><DEDENT>else:<EOL><INDENT>return _nested_convert_to_tensor(args, name=name)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>return nest.map_structure_up_to(<EOL>dtype, lambda s, dtype: _nested_convert_to_tensor(s, dtype, name), args,<EOL>dtype)<EOL><DEDENT>

Converts `args` to `Tensor`s.

    Use this when it is necessary to convert user-provided arguments that will
    then be passed to user-provided callables.

    When `dtype` is `None` this function behaves as follows:

    1A. If the top-level structure is a `list`/`tuple` but not a `namedtuple`,
        then it is left as is and only its elements are converted to `Tensor`s.

    2A. The sub-structures are converted to `Tensor`s eagerly. E.g. if `args` is
        `{'arg': [[1], [2]]}` it is converted to
        `{'arg': tf.constant([[1], [2]])}`. If the conversion fails, it will
        attempt to recurse into its children.

    When `dtype` is specified, it acts as both a structural and numeric type
    constraint. `dtype` can be a single `DType`, `None` or a nested collection
    thereof. The conversion rule becomes as follows:

    1B. The return value of this function will have the same structure as `dtype`.

    2B. If the leaf of `dtype` is a concrete `DType`, then the corresponding
        sub-structure in `args` is converted to a `Tensor`.

    3B. If the leaf of `dtype` is `None`, then the corresponding sub-structure is
        converted eagerly as described in the rule 2A above.

    Args:
      args: Arguments to convert to `Tensor`s.
      dtype: Optional structure/numeric type constraint.
      name: Optional name-scope to use.

    Returns:
      args: Converted `args`.

    #### Examples.

    This table shows some useful conversion cases. `T` means `Tensor`, `NT` means
    `namedtuple` and `CNT` means a `namedtuple` with a `Tensor`-conversion
    function registered.

    |     args     |    dtype   |       output       |
    |:------------:|:----------:|:------------------:|
    | `{"a": 1}`   | `None`     | `{"a": T(1)}`      |
    | `T(1)`       | `None`     | `T(1)`             |
    | `[1]`        | `None`     | `[T(1)]`           |
    | `[1]`        | `tf.int32` | `T([1])`           |
    | `[[T(1)]]`   | `None`     | `[T([1])]`         |
    | `[[T(1)]]`   | `[[None]]` | `[[T(1)]]`         |
    | `NT(1, 2)`   | `None`     | `NT(T(1), T(2))`   |
    | `NT(1, 2)`   | `tf.int32` | `T([1, 2])`        |
    | `CNT(1, 2)`  | `None`     | `T(...)`           |
    | `[[1, [2]]]` | `None`     | `[[T(1), T([2])]]` |


if expand_as_args(args):<EOL><INDENT>return fn(*args)<EOL><DEDENT>elif _expand_as_kwargs(args):<EOL><INDENT>return fn(**args)<EOL><DEDENT>else:<EOL><INDENT>return fn(args)<EOL><DEDENT>

Calls `fn` with `args`, possibly expanding `args`.

    Use this function when calling a user-provided callable using user-provided
    arguments.

    The expansion rules are as follows:

    `fn(*args)` if `args` is a `list` or a `tuple`, but not a `namedtuple`.
    `fn(**args)` if `args` is a `dict`.
    `fn(args)` otherwise.

    Args:
      fn: A callable that takes either `args` as an argument(s).
      args: Arguments to `fn`.

    Returns:
      result: Return value of `fn`.


def mixture_stddev(mixture_weight_vector, mean_vector, stddev_vector):

tensorshape_util.assert_has_rank(mixture_weight_vector.shape, <NUM_LIT:2>)<EOL>if not tensorshape_util.is_compatible_with(mean_vector.shape,<EOL>mixture_weight_vector.shape):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>if not tensorshape_util.is_compatible_with(stddev_vector.shape,<EOL>mixture_weight_vector.shape):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>pi_for_dot_prod = tf.expand_dims(mixture_weight_vector, axis=<NUM_LIT:1>)<EOL>mu_for_dot_prod = tf.expand_dims(mean_vector, axis=<NUM_LIT:2>)<EOL>sigma_for_dot_prod = tf.expand_dims(stddev_vector, axis=<NUM_LIT:2>)<EOL>mean_wa = tf.matmul(pi_for_dot_prod, mu_for_dot_prod)<EOL>mean_wa = tf.reshape(mean_wa, (-<NUM_LIT:1>,))<EOL>var_wa = tf.matmul(pi_for_dot_prod, tf.square(sigma_for_dot_prod))<EOL>var_wa = tf.reshape(var_wa, (-<NUM_LIT:1>,))<EOL>sq_mean_wa = tf.matmul(pi_for_dot_prod, tf.square(mu_for_dot_prod))<EOL>sq_mean_wa = tf.reshape(sq_mean_wa, (-<NUM_LIT:1>,))<EOL>mixture_variance = var_wa + sq_mean_wa - tf.square(mean_wa)<EOL>return tf.sqrt(mixture_variance)<EOL>

Computes the standard deviation of a mixture distribution.

    This function works regardless of the component distribution, so long as
    each component's mean and standard deviation can be provided.

    Args:
      mixture_weight_vector: A 2D tensor with shape [batch_size, num_components]
      mean_vector: A 2D tensor of mixture component means. Has shape `[batch_size,
        num_components]`.
      stddev_vector: A 2D tensor of mixture component standard deviations. Has
        shape `[batch_size, num_components]`.

    Returns:
      A 1D tensor of shape `[batch_size]` representing the standard deviation of
      the mixture distribution with given weights and component means and standard
      deviations.
    Raises:
      ValueError: If the shapes of the input tensors are not as expected.


def _maybe_attach_assertion(x):<EOL><INDENT>if not validate_args:<EOL><INDENT>return x<EOL><DEDENT>if assert_positive:<EOL><INDENT>return with_dependencies([<EOL>assert_util.assert_positive(<EOL>tf.linalg.diag_part(x), message="<STR_LIT>"),<EOL>], x)<EOL><DEDENT>return with_dependencies([<EOL>assert_util.assert_none_equal(<EOL>tf.linalg.diag_part(x),<EOL>tf.zeros([], x.dtype),<EOL>message="<STR_LIT>"),<EOL>], x)<EOL><DEDENT>with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[loc, scale_tril, scale_diag, scale_identity_multiplier],<EOL>preferred_dtype=tf.float32)<EOL>loc = _convert_to_tensor(loc, name="<STR_LIT>", dtype=dtype)<EOL>scale_tril = _convert_to_tensor(scale_tril, name="<STR_LIT>", dtype=dtype)<EOL>scale_diag = _convert_to_tensor(scale_diag, name="<STR_LIT>", dtype=dtype)<EOL>scale_identity_multiplier = _convert_to_tensor(<EOL>scale_identity_multiplier,<EOL>name="<STR_LIT>",<EOL>dtype=dtype)<EOL><DEDENT>if scale_tril is not None:<EOL><INDENT>scale_tril = tf.linalg.band_part(scale_tril, -<NUM_LIT:1>, <NUM_LIT:0>)  <EOL>tril_diag = tf.linalg.diag_part(scale_tril)<EOL>if scale_diag is not None:<EOL><INDENT>tril_diag += scale_diag<EOL><DEDENT>if scale_identity_multiplier is not None:<EOL><INDENT>tril_diag += scale_identity_multiplier[..., tf.newaxis]<EOL><DEDENT>scale_tril = tf.linalg.set_diag(scale_tril, tril_diag)<EOL>return tf.linalg.LinearOperatorLowerTriangular(<EOL>tril=_maybe_attach_assertion(scale_tril),<EOL>is_non_singular=True,<EOL>is_self_adjoint=False,<EOL>is_positive_definite=assert_positive)<EOL><DEDENT>return make_diag_scale(<EOL>loc=loc,<EOL>scale_diag=scale_diag,<EOL>scale_identity_multiplier=scale_identity_multiplier,<EOL>shape_hint=shape_hint,<EOL>validate_args=validate_args,<EOL>assert_positive=assert_positive,<EOL>name=name)<EOL>

Creates a LinearOperator representing a lower triangular matrix.

    Args:
      loc: Floating-point `Tensor`. This is used for inferring shape in the case
        where only `scale_identity_multiplier` is set.
      scale_tril: Floating-point `Tensor` representing the diagonal matrix.
        `scale_diag` has shape [N1, N2, ...  k, k], which represents a k x k lower
        triangular matrix. When `None` no `scale_tril` term is added to the
        LinearOperator. The upper triangular elements above the diagonal are
        ignored.
      scale_diag: Floating-point `Tensor` representing the diagonal matrix.
        `scale_diag` has shape [N1, N2, ...  k], which represents a k x k diagonal
        matrix. When `None` no diagonal term is added to the LinearOperator.
      scale_identity_multiplier: floating point rank 0 `Tensor` representing a
        scaling done to the identity matrix. When `scale_identity_multiplier =
        scale_diag = scale_tril = None` then `scale += IdentityMatrix`. Otherwise
        no scaled-identity-matrix is added to `scale`.
      shape_hint: scalar integer `Tensor` representing a hint at the dimension of
        the identity matrix when only `scale_identity_multiplier` is set.
      validate_args: Python `bool` indicating whether arguments should be checked
        for correctness.
      assert_positive: Python `bool` indicating whether LinearOperator should be
        checked for being positive definite.
      name: Python `str` name given to ops managed by this object.

    Returns:
      `LinearOperator` representing a lower triangular matrix.

    Raises:
      ValueError:  If only `scale_identity_multiplier` is set and `loc` and
        `shape_hint` are both None.


def _maybe_attach_assertion(x):<EOL><INDENT>if not validate_args:<EOL><INDENT>return x<EOL><DEDENT>if assert_positive:<EOL><INDENT>return with_dependencies([<EOL>assert_util.assert_positive(<EOL>x, message="<STR_LIT>"),<EOL>], x)<EOL><DEDENT>return with_dependencies([<EOL>assert_util.assert_none_equal(<EOL>x, tf.zeros([], x.dtype), message="<STR_LIT>")<EOL>], x)<EOL><DEDENT>with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>if dtype is None:<EOL><INDENT>dtype = dtype_util.common_dtype(<EOL>[loc, scale_diag, scale_identity_multiplier],<EOL>preferred_dtype=tf.float32)<EOL><DEDENT>loc = _convert_to_tensor(loc, name="<STR_LIT>", dtype=dtype)<EOL>scale_diag = _convert_to_tensor(scale_diag, name="<STR_LIT>", dtype=dtype)<EOL>scale_identity_multiplier = _convert_to_tensor(<EOL>scale_identity_multiplier,<EOL>name="<STR_LIT>",<EOL>dtype=dtype)<EOL>if scale_diag is not None:<EOL><INDENT>if scale_identity_multiplier is not None:<EOL><INDENT>scale_diag += scale_identity_multiplier[..., tf.newaxis]<EOL><DEDENT>return tf.linalg.LinearOperatorDiag(<EOL>diag=_maybe_attach_assertion(scale_diag),<EOL>is_non_singular=True,<EOL>is_self_adjoint=True,<EOL>is_positive_definite=assert_positive)<EOL><DEDENT>if loc is None and shape_hint is None:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>num_rows = shape_hint<EOL>del shape_hint<EOL>if num_rows is None:<EOL><INDENT>num_rows = tf.compat.dimension_value(loc.shape[-<NUM_LIT:1>])<EOL>if num_rows is None:<EOL><INDENT>num_rows = tf.shape(input=loc)[-<NUM_LIT:1>]<EOL><DEDENT><DEDENT>if scale_identity_multiplier is None:<EOL><INDENT>return tf.linalg.LinearOperatorIdentity(<EOL>num_rows=num_rows,<EOL>dtype=dtype,<EOL>is_self_adjoint=True,<EOL>is_positive_definite=True,<EOL>assert_proper_shapes=validate_args)<EOL><DEDENT>return tf.linalg.LinearOperatorScaledIdentity(<EOL>num_rows=num_rows,<EOL>multiplier=_maybe_attach_assertion(scale_identity_multiplier),<EOL>is_non_singular=True,<EOL>is_self_adjoint=True,<EOL>is_positive_definite=assert_positive,<EOL>assert_proper_shapes=validate_args)<EOL><DEDENT>

Creates a LinearOperator representing a diagonal matrix.

    Args:
      loc: Floating-point `Tensor`. This is used for inferring shape in the case
        where only `scale_identity_multiplier` is set.
      scale_diag: Floating-point `Tensor` representing the diagonal matrix.
        `scale_diag` has shape [N1, N2, ...  k], which represents a k x k diagonal
        matrix. When `None` no diagonal term is added to the LinearOperator.
      scale_identity_multiplier: floating point rank 0 `Tensor` representing a
        scaling done to the identity matrix. When `scale_identity_multiplier =
        scale_diag = scale_tril = None` then `scale += IdentityMatrix`. Otherwise
        no scaled-identity-matrix is added to `scale`.
      shape_hint: scalar integer `Tensor` representing a hint at the dimension of
        the identity matrix when only `scale_identity_multiplier` is set.
      validate_args: Python `bool` indicating whether arguments should be checked
        for correctness.
      assert_positive: Python `bool` indicating whether LinearOperator should be
        checked for being positive definite.
      name: Python `str` name given to ops managed by this object.
      dtype: TF `DType` to prefer when converting args to `Tensor`s. Else, we fall
        back to a compatible dtype across all of `loc`, `scale_diag`, and
        `scale_identity_multiplier`.

    Returns:
      `LinearOperator` representing a lower triangular matrix.

    Raises:
      ValueError:  If only `scale_identity_multiplier` is set and `loc` and
        `shape_hint` are both None.


def shapes_from_loc_and_scale(loc, scale, name="<STR_LIT>"):

if loc is not None and tensorshape_util.rank(loc.shape) == <NUM_LIT:0>:<EOL><INDENT>loc = None  <EOL><DEDENT>with tf.name_scope(name):<EOL><INDENT>event_size = scale.range_dimension_tensor()<EOL>event_size_ = tf.get_static_value(event_size)<EOL>loc_event_size_ = (None if loc is None<EOL>else tf.compat.dimension_value(loc.shape[-<NUM_LIT:1>]))<EOL>if event_size_ is not None and loc_event_size_ is not None:<EOL><INDENT>if loc_event_size_ != <NUM_LIT:1> and loc_event_size_ != event_size_:<EOL><INDENT>raise ValueError(<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(event_size_, loc_event_size_))<EOL><DEDENT><DEDENT>elif loc_event_size_ is not None and loc_event_size_ != <NUM_LIT:1>:<EOL><INDENT>event_size_ = loc_event_size_<EOL><DEDENT>if event_size_ is None:<EOL><INDENT>event_shape = event_size[tf.newaxis]<EOL><DEDENT>else:<EOL><INDENT>event_shape = tf.convert_to_tensor(<EOL>value=np.reshape(event_size_, [<NUM_LIT:1>]),<EOL>dtype=tf.int32,<EOL>name="<STR_LIT>")<EOL><DEDENT>batch_shape = scale.batch_shape_tensor()<EOL>if loc is not None:<EOL><INDENT>loc_batch_shape = tensorshape_util.with_rank_at_least(loc.shape, <NUM_LIT:1>)[:-<NUM_LIT:1>]<EOL>if tensorshape_util.rank(<EOL>loc.shape) is None or not tensorshape_util.is_fully_defined(<EOL>loc_batch_shape):<EOL><INDENT>loc_batch_shape = tf.shape(input=loc)[:-<NUM_LIT:1>]<EOL><DEDENT>else:<EOL><INDENT>loc_batch_shape = tf.convert_to_tensor(<EOL>value=loc_batch_shape, dtype=tf.int32, name="<STR_LIT>")<EOL><DEDENT>batch_shape = prefer_static_broadcast_shape(batch_shape, loc_batch_shape)  <EOL>batch_shape = tf.convert_to_tensor(<EOL>value=batch_shape, dtype=tf.int32, name="<STR_LIT>")<EOL><DEDENT>return batch_shape, event_shape<EOL><DEDENT>

Infer distribution batch and event shapes from a location and scale.

    Location and scale family distributions determine their batch/event shape by
    broadcasting the `loc` and `scale` args.  This helper does that broadcast,
    statically if possible.

    Batch shape broadcasts as per the normal rules.
    We allow the `loc` event shape to broadcast up to that of `scale`.  We do not
    allow `scale`'s event shape to change.  Therefore, the last dimension of `loc`
    must either be size `1`, or the same as `scale.range_dimension`.

    See `MultivariateNormalLinearOperator` for a usage example.

    Args:
      loc: `Tensor` (already converted to tensor) or `None`. If `None`, or
        `rank(loc)==0`, both batch and event shape are determined by `scale`.
      scale:  A `LinearOperator` instance.
      name:  A string name to prepend to created ops.

    Returns:
      batch_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.
      event_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.

    Raises:
      ValueError:  If the last dimension of `loc` is determined statically to be
        different than the range of `scale`.


<EOL>s_shape = tensors[<NUM_LIT:0>].shape<EOL>for t in tensors[<NUM_LIT:1>:]:<EOL><INDENT>s_shape = tf.broadcast_static_shape(s_shape, t.shape)<EOL><DEDENT>if tensorshape_util.is_fully_defined(s_shape):<EOL><INDENT>return tensorshape_util.as_list(s_shape)<EOL><DEDENT>d_shape = tf.shape(input=tensors[<NUM_LIT:0>])<EOL>for t in tensors[<NUM_LIT:1>:]:<EOL><INDENT>d_shape = tf.broadcast_dynamic_shape(d_shape, tf.shape(input=t))<EOL><DEDENT>return d_shape<EOL>

Get broadcast shape as a Python list of integers (preferred) or `Tensor`.

    Args:
      *tensors:  One or more `Tensor` objects (already converted!).

    Returns:
      broadcast shape:  Python list (if shapes determined statically), otherwise
        an `int32` `Tensor`.


if not isinstance(scale, tf.linalg.LinearOperator):<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>" % scale)<EOL><DEDENT>return (isinstance(scale, tf.linalg.LinearOperatorIdentity) or<EOL>isinstance(scale, tf.linalg.LinearOperatorScaledIdentity) or<EOL>isinstance(scale, tf.linalg.LinearOperatorDiag))<EOL>

Returns `True` if `scale` is a `LinearOperator` that is known to be diag.

    Args:
      scale:  `LinearOperator` instance.

    Returns:
      Python `bool`.

    Raises:
      TypeError:  If `scale` is not a `LinearOperator`.


if distribution.dtype != expected_base_dtype:<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>dtype_util.name(distribution.dtype),<EOL>dtype_util.name(expected_base_dtype)))<EOL><DEDENT>if validate_args and (distribution.reparameterization_type !=<EOL>reparameterization.FULLY_REPARAMETERIZED):<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>distribution.reparameterization_type))<EOL><DEDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>assertions = []<EOL>def check_is_scalar(is_scalar, name):<EOL><INDENT>is_scalar_ = tf.get_static_value(is_scalar)<EOL>if is_scalar_ is not None:<EOL><INDENT>if not is_scalar_:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>".format(name))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>assert_util.assert_equal(<EOL>is_scalar,<EOL>True,<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>".format(name))))<EOL><DEDENT><DEDENT>check_is_scalar(distribution.is_scalar_event(), "<STR_LIT>")<EOL>check_is_scalar(distribution.is_scalar_batch(), "<STR_LIT>")<EOL>return assertions<EOL><DEDENT>

Helper which checks validity of a scalar `distribution` init arg.

    Valid here means:

    * `distribution` has scalar batch and event shapes.
    * `distribution` is `FULLY_REPARAMETERIZED`
    * `distribution` has expected dtype.

    Args:
      distribution:  `Distribution`-like object.
      expected_base_dtype:  `TensorFlow` `dtype`.
      validate_args:  Python `bool`.  Whether to do additional checks: (i)  check
        that reparameterization_type is `FULLY_REPARAMETERIZED`. (ii) add
        `tf.Assert` ops to the graph to enforce that distribution is scalar in the
        event that this cannot be determined statically.

    Returns:
      List of `tf.Assert` ops to run to enforce validity checks that could not
        be statically determined.  Empty if `not validate_args`.

    Raises:
      ValueError:  If validate_args and distribution is not FULLY_REPARAMETERIZED
      ValueError:  If distribution is statically determined to not have both
        scalar batch and scalar event shapes.


def pad_mixture_dimensions(x, mixture_distribution, categorical_distribution,<EOL>event_ndims):

with tf.name_scope("<STR_LIT>"):<EOL><INDENT>def _get_ndims(d):<EOL><INDENT>if tensorshape_util.rank(d.batch_shape) is not None:<EOL><INDENT>return tensorshape_util.rank(d.batch_shape)<EOL><DEDENT>return tf.shape(input=d.batch_shape_tensor())[<NUM_LIT:0>]<EOL><DEDENT>dist_batch_ndims = _get_ndims(mixture_distribution)<EOL>cat_batch_ndims = _get_ndims(categorical_distribution)<EOL>pad_ndims = tf.where(categorical_distribution.is_scalar_batch(),<EOL>dist_batch_ndims, dist_batch_ndims - cat_batch_ndims)<EOL>s = tf.shape(input=x)<EOL>x = tf.reshape(<EOL>x,<EOL>shape=tf.concat([<EOL>s[:-<NUM_LIT:1>],<EOL>tf.ones([pad_ndims], dtype=tf.int32),<EOL>s[-<NUM_LIT:1>:],<EOL>tf.ones([event_ndims], dtype=tf.int32),<EOL>],<EOL>axis=<NUM_LIT:0>))<EOL>return x<EOL><DEDENT>

Pad dimensions of event tensors for mixture distributions.

    See `Mixture._sample_n` and `MixtureSameFamily._sample_n` for usage examples.

    Args:
      x: event tensor to pad.
      mixture_distribution: Base distribution of the mixture.
      categorical_distribution: `Categorical` distribution that mixes the base
        distribution.
      event_ndims: Integer specifying the number of event dimensions in the event
        tensor.

    Returns:
      A padded version of `x` that can broadcast with `categorical_distribution`.


def pick_scalar_condition(pred, true_value, false_value, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>pred = tf.convert_to_tensor(<EOL>value=pred, dtype_hint=tf.bool, name="<STR_LIT>")<EOL>true_value = tf.convert_to_tensor(value=true_value, name="<STR_LIT>")<EOL>false_value = tf.convert_to_tensor(value=false_value, name="<STR_LIT>")<EOL>pred_ = tf.get_static_value(pred)<EOL>if pred_ is None:<EOL><INDENT>return tf.where(pred, true_value, false_value)<EOL><DEDENT>return true_value if pred_ else false_value<EOL><DEDENT>

Convenience function that chooses one of two values based on the predicate.

    This utility is equivalent to a version of `tf.where` that accepts only a
    scalar predicate and computes its result statically when possible. It may also
    be used in place of `tf.cond` when both branches yield a `Tensor` of the same
    shape; the operational difference is that `tf.cond` uses control flow to
    evaluate only the branch that's needed, while `tf.where` (and thus
    this method) may evaluate both branches before the predicate's truth is known.
    This means that `tf.cond` is preferred when one of the branches is expensive
    to evaluate (like performing a large matmul), while this method is preferred
    when both branches are cheap, e.g., constants. In the latter case, we expect
    this method to be substantially faster than `tf.cond` on GPU and to give
    similar performance on CPU.

    Args:
      pred: Scalar `bool` `Tensor` predicate.
      true_value: `Tensor` to return if `pred` is `True`.
      false_value: `Tensor` to return if `pred` is `False`. Must have the same
        shape as `true_value`.
      name: Python `str` name given to ops managed by this object.

    Returns:
      result: a `Tensor` (or `Tensor`-convertible Python value) equal to
        `true_value` if `pred` evaluates to `True` and `false_value` otherwise.
        If the condition can be evaluated statically, the result returned is one
        of the input Python values, with no graph side effects.


axis = tf.convert_to_tensor(value=axis, name="<STR_LIT>")<EOL>rank = tf.convert_to_tensor(value=rank, name="<STR_LIT>")<EOL>axis_ = tf.get_static_value(axis)<EOL>rank_ = tf.get_static_value(rank)<EOL>if axis_ is not None and rank_ is not None:<EOL><INDENT>is_scalar = axis_.ndim == <NUM_LIT:0><EOL>if is_scalar:<EOL><INDENT>axis_ = [axis_]<EOL><DEDENT>positive_axis = []<EOL>for a_ in axis_:<EOL><INDENT>if a_ < <NUM_LIT:0>:<EOL><INDENT>positive_axis.append(rank_ + a_)<EOL><DEDENT>else:<EOL><INDENT>positive_axis.append(a_)<EOL><DEDENT><DEDENT>if is_scalar:<EOL><INDENT>positive_axis = positive_axis[<NUM_LIT:0>]<EOL><DEDENT>return tf.convert_to_tensor(value=positive_axis, dtype=axis.dtype)<EOL><DEDENT>return tf.where(axis < <NUM_LIT:0>, rank + axis, axis)<EOL>

Make (possibly negatively indexed) `axis` argument non-negative.


def move_dimension(x, source_idx, dest_idx):

ndims = prefer_static_rank(x)<EOL>dtype = dtype_util.common_dtype([source_idx, dest_idx],<EOL>preferred_dtype=tf.int32)<EOL>source_idx = tf.convert_to_tensor(value=source_idx, dtype=dtype)<EOL>dest_idx = tf.convert_to_tensor(value=dest_idx, dtype=dtype)<EOL>source_idx = pick_scalar_condition(source_idx < <NUM_LIT:0>, ndims + source_idx,<EOL>source_idx)<EOL>dest_idx = pick_scalar_condition(dest_idx < <NUM_LIT:0>, ndims + dest_idx, dest_idx)<EOL>def move_left_permutation():<EOL><INDENT>return prefer_static_value(<EOL>tf.concat([<EOL>tf.range(<NUM_LIT:0>, dest_idx, dtype=dtype), [source_idx],<EOL>tf.range(dest_idx, source_idx, dtype=dtype),<EOL>tf.range(source_idx + <NUM_LIT:1>, ndims, dtype=dtype)<EOL>],<EOL>axis=<NUM_LIT:0>))<EOL><DEDENT>def move_right_permutation():<EOL><INDENT>return prefer_static_value(<EOL>tf.concat([<EOL>tf.range(<NUM_LIT:0>, source_idx, dtype=dtype),<EOL>tf.range(source_idx + <NUM_LIT:1>, dest_idx + <NUM_LIT:1>, dtype=dtype), [source_idx],<EOL>tf.range(dest_idx + <NUM_LIT:1>, ndims, dtype=dtype)<EOL>],<EOL>axis=<NUM_LIT:0>))<EOL><DEDENT>def x_permuted():<EOL><INDENT>return tf.transpose(<EOL>a=x,<EOL>perm=prefer_static.cond(source_idx < dest_idx,<EOL>move_right_permutation,<EOL>move_left_permutation))<EOL><DEDENT>return prefer_static.cond(tf.equal(source_idx, dest_idx),<EOL>lambda: x, x_permuted)<EOL>

Move a single tensor dimension within its shape.

    This is a special case of `tf.transpose()`, which applies
    arbitrary permutations to tensor dimensions.

    Args:
      x: Tensor of rank `ndims`.
      source_idx: Integer index into `x.shape` (negative indexing is supported).
      dest_idx: Integer index into `x.shape` (negative indexing is supported).

    Returns:
      x_perm: Tensor of rank `ndims`, in which the dimension at original
       index `source_idx` has been moved to new index `dest_idx`, with
       all other dimensions retained in their original order.

    Example:

    ```python
    x = tf.placeholder(shape=[200, 30, 4, 1, 6])
    x_perm = _move_dimension(x, 1, 1) # no-op
    x_perm = _move_dimension(x, 0, 3) # result shape [30, 4, 1, 200, 6]
    x_perm = _move_dimension(x, 0, -2) # equivalent to previous
    x_perm = _move_dimension(x, 4, 2) # result shape [200, 30, 6, 4, 1]
    ```


with tf.name_scope(name):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>if dtype_util.is_integer(x.dtype):<EOL><INDENT>return tf.no_op()<EOL><DEDENT>message = message or "<STR_LIT>".format(x)<EOL>if int_dtype is None:<EOL><INDENT>try:<EOL><INDENT>int_dtype = {<EOL>tf.float16: tf.int16,<EOL>tf.float32: tf.int32,<EOL>tf.float64: tf.int64,<EOL>}[dtype_util.base_dtype(x.dtype)]<EOL><DEDENT>except KeyError:<EOL><INDENT>raise TypeError("<STR_LIT>".format(dtype_util.name(x.dtype)))<EOL><DEDENT><DEDENT>return assert_util.assert_equal(<EOL>x,<EOL>tf.cast(tf.cast(x, int_dtype), x.dtype),<EOL>data=data,<EOL>summarize=summarize,<EOL>message=message,<EOL>name=name)<EOL><DEDENT>

Assert that x has integer components (or floats equal to integers).

    Args:
      x: Floating-point `Tensor`
      data: The tensors to print out if the condition is `False`. Defaults to
        error message and first few entries of `x` and `y`.
      summarize: Print this many entries of each tensor.
      message: A string to prefix to the default message.
      int_dtype: A `tf.dtype` used to cast the float to. The default (`None`)
        implies the smallest possible signed int will be used for casting.
      name: A name for this operation (optional).

    Returns:
      Op raising `InvalidArgumentError` if `cast(x, int_dtype) != x`.


with tf.name_scope(name):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>assertions = [<EOL>assert_util.assert_non_negative(<EOL>x, message="<STR_LIT>".format(x)),<EOL>]<EOL>if not dtype_util.is_integer(x.dtype):<EOL><INDENT>assertions += [<EOL>assert_integer_form(<EOL>x,<EOL>message="<STR_LIT>".format(x)),<EOL>]<EOL><DEDENT>return with_dependencies(assertions, x)<EOL><DEDENT>

Assert x is a non-negative tensor, and optionally of integers.


a = tf.convert_to_tensor(value=a, name="<STR_LIT:a>")<EOL>b = tf.convert_to_tensor(value=b, name="<STR_LIT:b>")<EOL>def all_shapes_equal():<EOL><INDENT>return tf.reduce_all(<EOL>input_tensor=tf.equal(<EOL>tf.concat([tf.shape(input=a), tf.shape(input=b)], <NUM_LIT:0>),<EOL>tf.concat([tf.shape(input=b), tf.shape(input=a)], <NUM_LIT:0>)))<EOL><DEDENT>return tf.cond(<EOL>pred=tf.equal(tf.rank(a), tf.rank(b)),<EOL>true_fn=all_shapes_equal,<EOL>false_fn=lambda: tf.constant(False))<EOL>

Returns whether a and b have the same dynamic shape.

    Args:
      a: `Tensor`
      b: `Tensor`

    Returns:
      `bool` `Tensor` representing if both tensors have the same shape.


if x is None:<EOL><INDENT>return x<EOL><DEDENT>try:<EOL><INDENT>x_ = tf.get_static_value(x)<EOL><DEDENT>except TypeError:<EOL><INDENT>x_ = x<EOL><DEDENT>if x_ is None or dtype is None:<EOL><INDENT>return x_<EOL><DEDENT>return np.array(x_, dtype)<EOL>

Helper which tries to return a static value.

    Given `x`, extract it's value statically, optionally casting to a specific
    dtype. If this is not possible, None is returned.

    Args:
      x: `Tensor` for which to extract a value statically.
      dtype: Optional dtype to cast to.

    Returns:
      Statically inferred value if possible, otherwise None.


if dtype is None:<EOL><INDENT>dtype = dtype_util.common_dtype([probs, logits], preferred_dtype=tf.float32)<EOL><DEDENT>with tf.name_scope(name):<EOL><INDENT>if (probs is None) == (logits is None):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>if probs is None:<EOL><INDENT>logits = tf.convert_to_tensor(value=logits, name="<STR_LIT>", dtype=dtype)<EOL>if not dtype_util.is_floating(logits.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>if multidimensional:<EOL><INDENT>if validate_args:<EOL><INDENT>logits = embed_check_categorical_event_shape(logits)<EOL><DEDENT>return logits, tf.nn.softmax(logits, name="<STR_LIT>")<EOL><DEDENT>return logits, tf.sigmoid(logits, name="<STR_LIT>")<EOL><DEDENT>probs = tf.convert_to_tensor(value=probs, name="<STR_LIT>", dtype=dtype)<EOL>if not dtype_util.is_floating(probs.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>if validate_args:<EOL><INDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>one = tf.constant(<NUM_LIT:1.>, probs.dtype)<EOL>dependencies = [assert_util.assert_non_negative(probs)]<EOL>if multidimensional:<EOL><INDENT>probs = embed_check_categorical_event_shape(probs)<EOL>dependencies += [<EOL>assert_util.assert_near(<EOL>tf.reduce_sum(input_tensor=probs, axis=-<NUM_LIT:1>),<EOL>one,<EOL>message="<STR_LIT>")<EOL>]<EOL><DEDENT>else:<EOL><INDENT>dependencies += [<EOL>assert_util.assert_less_equal(<EOL>probs, one, message="<STR_LIT>")<EOL>]<EOL><DEDENT>probs = with_dependencies(dependencies, probs)<EOL><DEDENT><DEDENT>with tf.name_scope("<STR_LIT>"):<EOL><INDENT>if multidimensional:<EOL><INDENT>return tf.math.log(probs), probs<EOL><DEDENT>return tf.math.log(probs) - tf.math.log1p(-<NUM_LIT:1.> * probs), probs<EOL><DEDENT><DEDENT>

Converts logit to probabilities (or vice-versa), and returns both.

    Args:
      logits: Floating-point `Tensor` representing log-odds.
      probs: Floating-point `Tensor` representing probabilities.
      multidimensional: Python `bool`, default `False`. If `True`, represents
        whether the last dimension of `logits` or `probs`, a `[N1, N2, ...  k]`
        dimensional tensor, representing the logit or probability of `shape[-1]`
        classes.
      validate_args: Python `bool`, default `False`. When `True`, either assert `0
        <= probs <= 1` (if not `multidimensional`) or that the last dimension of
        `probs` sums to one.
      name: A name for this operation (optional).
      dtype: `tf.DType` to prefer when converting args to `Tensor`s.

    Returns:
      logits, probs: Tuple of `Tensor`s. If `probs` has an entry that is `0` or
        `1`, then the corresponding entry in the returned logit will be `-Inf` and
        `Inf` respectively.

    Raises:
      ValueError: if neither `probs` nor `logits` were passed in, or both were.


return {<EOL>tf.bool: True,<EOL>tf.uint8: True,<EOL>tf.uint16: True,<EOL>}.get(dt.base_dtype, False)<EOL>

Helper returning True if dtype is known to be unsigned.


return {<EOL>tf.float16: True,<EOL>tf.float32: True,<EOL>tf.float64: True,<EOL>tf.int8: True,<EOL>tf.int16: True,<EOL>tf.int32: True,<EOL>tf.int64: True,<EOL>}.get(dt.base_dtype, False)<EOL>

Helper returning True if dtype is known to be signed.


return _is_known_unsigned_by_dtype(dt) or _is_known_signed_by_dtype(dt)<EOL>

Helper returning True if dtype is known.


if not _is_known_dtype(dt):<EOL><INDENT>raise TypeError("<STR_LIT>".format(dt.name))<EOL><DEDENT>if dt.is_floating:<EOL><INDENT>return int(<NUM_LIT:2>**(np.finfo(dt.as_numpy_dtype).nmant + <NUM_LIT:1>))<EOL><DEDENT>if dt.is_integer:<EOL><INDENT>return np.iinfo(dt.as_numpy_dtype).max<EOL><DEDENT>if dt.base_dtype == tf.bool:<EOL><INDENT>return int(<NUM_LIT:1>)<EOL><DEDENT>raise TypeError("<STR_LIT>".format(dt.name))<EOL>

Helper returning the largest integer exactly representable by dtype.


if not _is_known_dtype(dt):<EOL><INDENT>raise TypeError("<STR_LIT>".format(dt.name))<EOL><DEDENT>if _is_known_unsigned_by_dtype(dt):<EOL><INDENT>return <NUM_LIT:0><EOL><DEDENT>return -<NUM_LIT:1> * _largest_integer_by_dtype(dt)<EOL>

Helper returning the smallest integer exactly representable by dtype.


if not _is_known_dtype(dt):<EOL><INDENT>raise TypeError("<STR_LIT>".format(dt.name))<EOL><DEDENT>return dt.is_integer or dt.base_dtype == tf.bool<EOL>

Helper returning True if dtype.is_integer or is `bool`.


with tf.name_scope(name):<EOL><INDENT>x = tf.convert_to_tensor(value=categorical_param, name="<STR_LIT>")<EOL>x_dtype = dtype_util.base_dtype(x.dtype)<EOL>max_event_size = (<EOL>_largest_integer_by_dtype(x_dtype)<EOL>if dtype_util.is_floating(x_dtype) else <NUM_LIT:0>)<EOL>if max_event_size is <NUM_LIT:0>:<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(dtype_util.name(x_dtype)))<EOL><DEDENT>try:<EOL><INDENT>x_shape_static = tensorshape_util.with_rank_at_least(x.shape, <NUM_LIT:1>)<EOL><DEDENT>except ValueError:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>event_size = tf.compat.dimension_value(x_shape_static[-<NUM_LIT:1>])<EOL>if event_size is not None:<EOL><INDENT>if event_size < <NUM_LIT:2>:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>if event_size > max_event_size:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>dtype_util.name(x_dtype), event_size,<EOL>max_event_size))<EOL><DEDENT>return x<EOL><DEDENT>else:<EOL><INDENT>event_size = tf.shape(input=x, out_type=tf.int64, name="<STR_LIT>")[-<NUM_LIT:1>]<EOL>return with_dependencies([<EOL>assert_util.assert_rank_at_least(<EOL>x,<EOL><NUM_LIT:1>,<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>")),<EOL>assert_util.assert_greater_equal(<EOL>tf.shape(input=x)[-<NUM_LIT:1>],<EOL><NUM_LIT:2>,<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>")),<EOL>assert_util.assert_less_equal(<EOL>event_size,<EOL>tf.convert_to_tensor(max_event_size, dtype=tf.int64),<EOL>message="<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>dtype_util.name(x_dtype), max_event_size)),<EOL>], x)<EOL><DEDENT><DEDENT>

Embeds checks that categorical distributions don't have too many classes.

    A categorical-type distribution is one which, e.g., returns the class label
    rather than a one-hot encoding.  E.g., `Categorical(probs)`.

    Since distributions output samples in the same dtype as the parameters, we
    must ensure that casting doesn't lose precision. That is, the
    `parameter.dtype` implies a maximum number of classes. However, since shape is
    `int32` and categorical variables are presumed to be indexes into a `Tensor`,
    we must also ensure that the number of classes is no larger than the largest
    possible `int32` index, i.e., `2**31-1`.

    In other words the number of classes, `K`, must satisfy the following
    condition:

    ```python
    K <= min(
        int(2**31 - 1),  # Largest float as an index.
        {
            tf.float16: int(2**11),   # Largest int as a float16.
            tf.float32: int(2**24),
            tf.float64: int(2**53),
        }.get(dtype_util.base_dtype(categorical_param.dtype), 0))
    ```

    Args:
      categorical_param: Floating-point `Tensor` representing parameters of
        distribution over categories. The rightmost shape is presumed to be the
        number of categories.
      name: A name for this operation (optional).

    Returns:
      categorical_param: Input `Tensor` with appropriate assertions embedded.

    Raises:
      TypeError: if `categorical_param` has an unknown `dtype`.
      ValueError: if we can statically identify `categorical_param` as being too
        large (for being closed under int32/float casting).


with tf.name_scope(name):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>if (not _is_integer_like_by_dtype(x.dtype) and<EOL>not dtype_util.is_floating(x.dtype)):<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(dtype_util.name(x.dtype)))<EOL><DEDENT>if (not _is_integer_like_by_dtype(target_dtype) and<EOL>not dtype_util.is_floating(target_dtype)):<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(dtype_util.name(target_dtype)))<EOL><DEDENT>if (not _is_integer_like_by_dtype(x.dtype) and<EOL>not _is_integer_like_by_dtype(target_dtype)):<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>x, dtype_util.name(x.dtype),<EOL>dtype_util.name(target_dtype)))<EOL><DEDENT>assertions = []<EOL>if assert_positive:<EOL><INDENT>assertions += [<EOL>assert_util.assert_positive(x, message="<STR_LIT>"),<EOL>]<EOL><DEDENT>elif assert_nonnegative:<EOL><INDENT>assertions += [<EOL>assert_util.assert_non_negative(<EOL>x, message="<STR_LIT>"),<EOL>]<EOL><DEDENT>if dtype_util.is_floating(x.dtype):<EOL><INDENT>assertions += [<EOL>assert_integer_form(<EOL>x,<EOL>int_dtype=target_dtype,<EOL>message="<STR_LIT>".format(<EOL>dtype_util.name(target_dtype))),<EOL>]<EOL><DEDENT>else:<EOL><INDENT>if (_largest_integer_by_dtype(x.dtype) ><EOL>_largest_integer_by_dtype(target_dtype)):<EOL><INDENT>assertions += [<EOL>assert_util.assert_less_equal(<EOL>x,<EOL>_largest_integer_by_dtype(target_dtype),<EOL>message=("<STR_LIT>".format(<EOL>_largest_integer_by_dtype(target_dtype)))),<EOL>]<EOL><DEDENT>if (not assert_nonnegative and (_smallest_integer_by_dtype(<EOL>x.dtype) < _smallest_integer_by_dtype(target_dtype))):<EOL><INDENT>assertions += [<EOL>assert_util.assert_greater_equal(<EOL>x,<EOL>_smallest_integer_by_dtype(target_dtype),<EOL>message=("<STR_LIT>".format(<EOL>_smallest_integer_by_dtype(target_dtype)))),<EOL>]<EOL><DEDENT><DEDENT>if not assertions:<EOL><INDENT>return x<EOL><DEDENT>return with_dependencies(assertions, x)<EOL><DEDENT>

Ensures integers remain unaffected despite casting to/from int/float types.

    Example integer-types: `uint8`, `int32`, `bool`.
    Example floating-types: `float32`, `float64`.

    The largest possible integer representable by an IEEE754 floating-point is
    `2**(1 + mantissa_bits)` yet the largest possible integer as an int-type is
    `2**(bits - 1) - 1`. This function ensures that a `Tensor` purporting to have
    integer-form values can be cast to some other type without loss of precision.

    The smallest representable integer is the negative of the largest
    representable integer, except for types: `uint8`, `uint16`, `bool`. For these
    types, the smallest representable integer is `0`.

    Args:
      x: `Tensor` representing integer-form values.
      target_dtype: TF `dtype` under which `x` should have identical values.
      assert_nonnegative: `bool` indicating `x` should contain nonnegative values.
      assert_positive: `bool` indicating `x` should contain positive values.
      name: A name for this operation (optional).

    Returns:
      x: Input `Tensor` with appropriate assertions embedded.

    Raises:
      TypeError: if `x` is neither integer- nor floating-type.
      TypeError: if `target_dtype` is neither integer- nor floating-type.
      TypeError: if neither `x` nor `target_dtype` are integer-type.


def log_combinations(n, counts, name="<STR_LIT>"):

<EOL>with tf.name_scope(name):<EOL><INDENT>n = tf.convert_to_tensor(value=n, name="<STR_LIT:n>")<EOL>counts = tf.convert_to_tensor(value=counts, name="<STR_LIT>")<EOL>total_permutations = tf.math.lgamma(n + <NUM_LIT:1>)<EOL>counts_factorial = tf.math.lgamma(counts + <NUM_LIT:1>)<EOL>redundant_permutations = tf.reduce_sum(<EOL>input_tensor=counts_factorial, axis=[-<NUM_LIT:1>])<EOL>return total_permutations - redundant_permutations<EOL><DEDENT>

Multinomial coefficient.

    Given `n` and `counts`, where `counts` has last dimension `k`, we compute
    the multinomial coefficient as:

    ```n! / sum_i n_i!```

    where `i` runs over all `k` classes.

    Args:
      n: Floating-point `Tensor` broadcastable with `counts`. This represents `n`
        outcomes.
      counts: Floating-point `Tensor` broadcastable with `n`. This represents
        counts in `k` classes, where `k` is the last dimension of the tensor.
      name: A name for this operation (optional).

    Returns:
      `Tensor` representing the multinomial coefficient between `n` and `counts`.


def matrix_diag_transform(matrix, transform=None, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>matrix = tf.convert_to_tensor(value=matrix, name="<STR_LIT>")<EOL>if transform is None:<EOL><INDENT>return matrix<EOL><DEDENT>diag = tf.linalg.diag_part(matrix)<EOL>transformed_diag = transform(diag)<EOL>transformed_mat = tf.linalg.set_diag(matrix, transformed_diag)<EOL><DEDENT>return transformed_mat<EOL>

Transform diagonal of [batch-]matrix, leave rest of matrix unchanged.

    Create a trainable covariance defined by a Cholesky factor:

    ```python
    # Transform network layer into 2 x 2 array.
    matrix_values = tf.contrib.layers.fully_connected(activations, 4)
    matrix = tf.reshape(matrix_values, (batch_size, 2, 2))

    # Make the diagonal positive. If the upper triangle was zero, this would be a
    # valid Cholesky factor.
    chol = matrix_diag_transform(matrix, transform=tf.nn.softplus)

    # LinearOperatorLowerTriangular ignores the upper triangle.
    operator = LinearOperatorLowerTriangular(chol)
    ```

    Example of heteroskedastic 2-D linear regression.

    ```python
    tfd = tfp.distributions

    # Get a trainable Cholesky factor.
    matrix_values = tf.contrib.layers.fully_connected(activations, 4)
    matrix = tf.reshape(matrix_values, (batch_size, 2, 2))
    chol = matrix_diag_transform(matrix, transform=tf.nn.softplus)

    # Get a trainable mean.
    mu = tf.contrib.layers.fully_connected(activations, 2)

    # This is a fully trainable multivariate normal!
    dist = tfd.MultivariateNormalTriL(mu, chol)

    # Standard log loss. Minimizing this will "train" mu and chol, and then dist
    # will be a distribution predicting labels as multivariate Gaussians.
    loss = -1 * tf.reduce_mean(dist.log_prob(labels))
    ```

    Args:
      matrix:  Rank `R` `Tensor`, `R >= 2`, where the last two dimensions are
        equal.
      transform:  Element-wise function mapping `Tensors` to `Tensors`. To be
        applied to the diagonal of `matrix`. If `None`, `matrix` is returned
        unchanged. Defaults to `None`.
      name:  A name to give created ops. Defaults to "matrix_diag_transform".

    Returns:
      A `Tensor` with same shape and `dtype` as `matrix`.


def rotate_transpose(x, shift, name="<STR_LIT>"):

with tf.name_scope(name):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>shift = tf.convert_to_tensor(value=shift, name="<STR_LIT>")<EOL>assert_util.assert_integer(shift)<EOL>shift_value_static = tf.get_static_value(shift)<EOL>ndims = tensorshape_util.rank(x.shape)<EOL>if ndims is not None and shift_value_static is not None:<EOL><INDENT>if ndims < <NUM_LIT:2>:<EOL><INDENT>return x<EOL><DEDENT>shift_value_static = np.sign(shift_value_static) * (<EOL>abs(shift_value_static) % ndims)<EOL>if shift_value_static == <NUM_LIT:0>:<EOL><INDENT>return x<EOL><DEDENT>perm = np.roll(np.arange(ndims), shift_value_static)<EOL>return tf.transpose(a=x, perm=perm)<EOL><DEDENT>else:<EOL><INDENT>ndims = tf.rank(x)<EOL>shift = tf.where(<EOL>tf.less(shift, <NUM_LIT:0>), -shift % ndims,<EOL>ndims - shift % ndims)<EOL>first = tf.range(<NUM_LIT:0>, shift)<EOL>last = tf.range(shift, ndims)<EOL>perm = tf.concat([last, first], <NUM_LIT:0>)<EOL>return tf.transpose(a=x, perm=perm)<EOL><DEDENT><DEDENT>

Circularly moves dims left or right.

    Effectively identical to:

    ```python
    numpy.transpose(x, numpy.roll(numpy.arange(len(x.shape)), shift))
    ```

    When `validate_args=False` additional graph-runtime checks are
    performed. These checks entail moving data from to GPU to CPU.

    Example:

    ```python
    x = tf.random_normal([1, 2, 3, 4])  # Tensor of shape [1, 2, 3, 4].
    rotate_transpose(x, -1).shape == [2, 3, 4, 1]
    rotate_transpose(x, -2).shape == [3, 4, 1, 2]
    rotate_transpose(x,  1).shape == [4, 1, 2, 3]
    rotate_transpose(x,  2).shape == [3, 4, 1, 2]
    rotate_transpose(x,  7).shape == rotate_transpose(x, 3).shape  # [2, 3, 4, 1]
    rotate_transpose(x, -7).shape == rotate_transpose(x, -3).shape  # [4, 1, 2, 3]
    ```

    Args:
      x: `Tensor`.
      shift: `Tensor`. Number of dimensions to transpose left (shift<0) or
        transpose right (shift>0).
      name: Python `str`. The name to give this op.

    Returns:
      rotated_x: Input `Tensor` with dimensions circularly rotated by shift.

    Raises:
      TypeError: if shift is not integer type.


def pick_vector(cond, true_vector, false_vector, name="<STR_LIT>"):

with tf.name_scope(name):<EOL><INDENT>cond = tf.convert_to_tensor(<EOL>value=cond, dtype_hint=tf.bool, name="<STR_LIT>")<EOL>if cond.dtype != tf.bool:<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>".format(cond, cond.dtype, tf.bool))<EOL><DEDENT>true_vector = tf.convert_to_tensor(value=true_vector, name="<STR_LIT>")<EOL>false_vector = tf.convert_to_tensor(value=false_vector, name="<STR_LIT>")<EOL>if true_vector.dtype != false_vector.dtype:<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>".format(<EOL>true_vector, true_vector.dtype, false_vector, false_vector.dtype))<EOL><DEDENT>cond_value_static = tf.get_static_value(cond)<EOL>if cond_value_static is not None:<EOL><INDENT>return true_vector if cond_value_static else false_vector<EOL><DEDENT>n = tf.shape(input=true_vector)[<NUM_LIT:0>]<EOL>return tf.slice(<EOL>tf.concat([true_vector, false_vector], <NUM_LIT:0>), [tf.where(cond, <NUM_LIT:0>, n)],<EOL>[tf.where(cond, n, -<NUM_LIT:1>)])<EOL><DEDENT>

Picks possibly different length row `Tensor`s based on condition.

    Value `Tensor`s should have exactly one dimension.

    If `cond` is a python Boolean or `tf.constant` then either `true_vector` or
    `false_vector` is immediately returned. I.e., no graph nodes are created and
    no validation happens.

    Args:
      cond: `Tensor`. Must have `dtype=tf.bool` and be scalar.
      true_vector: `Tensor` of one dimension. Returned when cond is `True`.
      false_vector: `Tensor` of one dimension. Returned when cond is `False`.
      name: Python `str`. The name to give this op.
    Example:  ```python pick_vector(tf.less(0, 5), tf.range(10, 12), tf.range(15,
      18))  # [10, 11] pick_vector(tf.less(5, 0), tf.range(10, 12), tf.range(15,
      18))  # [15, 16, 17] ```

    Returns:
      true_or_false_vector: `Tensor`.

    Raises:
      TypeError: if `cond.dtype != tf.bool`
      TypeError: if `cond` is not a constant and
        `true_vector.dtype != false_vector.dtype`


with tf.name_scope(name):<EOL><INDENT>def make_shape_tensor(x):<EOL><INDENT>return tf.convert_to_tensor(value=x, name="<STR_LIT>", dtype=tf.int32)<EOL><DEDENT>def get_tensor_shape(s):<EOL><INDENT>if isinstance(s, tf.TensorShape):<EOL><INDENT>return s<EOL><DEDENT>s_ = tf.get_static_value(make_shape_tensor(s))<EOL>if s_ is not None:<EOL><INDENT>return tf.TensorShape(s_)<EOL><DEDENT>return None<EOL><DEDENT>def get_shape_tensor(s):<EOL><INDENT>if not isinstance(s, tf.TensorShape):<EOL><INDENT>return make_shape_tensor(s)<EOL><DEDENT>if tensorshape_util.is_fully_defined(s):<EOL><INDENT>return make_shape_tensor(tensorshape_util.as_list(s))<EOL><DEDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>shape1_ = get_tensor_shape(shape1)<EOL>shape2_ = get_tensor_shape(shape2)<EOL>if shape1_ is not None and shape2_ is not None:<EOL><INDENT>return tf.broadcast_static_shape(shape1_, shape2_)<EOL><DEDENT>shape1_ = get_shape_tensor(shape1)<EOL>shape2_ = get_shape_tensor(shape2)<EOL>return tf.broadcast_dynamic_shape(shape1_, shape2_)<EOL><DEDENT>

Convenience function which statically broadcasts shape when possible.

    Args:
      shape1:  `1-D` integer `Tensor`.  Already converted to tensor!
      shape2:  `1-D` integer `Tensor`.  Already converted to tensor!
      name:  A string name to prepend to created ops.

    Returns:
      The broadcast shape, either as `TensorShape` (if broadcast can be done
        statically), or as a `Tensor`.


return prefer_static_value(tf.rank(x))<EOL>

Return static rank of tensor `x` if available, else `tf.rank(x)`.

    Args:
      x: `Tensor` (already converted).

    Returns:
      Numpy array (if static rank is obtainable), else `Tensor`.


return prefer_static_value(tf.shape(input=x))<EOL>

Return static shape of tensor `x` if available, else `tf.shape(x)`.

    Args:
      x: `Tensor` (already converted).

    Returns:
      Numpy array (if static shape is obtainable), else `Tensor`.


static_x = tf.get_static_value(x)<EOL>if static_x is not None:<EOL><INDENT>return static_x<EOL><DEDENT>return x<EOL>

Return static value of tensor `x` if available, else `x`.

    Args:
      x: `Tensor` (already converted).

    Returns:
      Numpy array (if static value is obtainable), else `Tensor`.


if seed is None:<EOL><INDENT>return None<EOL><DEDENT>string = (str(seed) + salt).encode("<STR_LIT:utf-8>")<EOL>return int(hashlib.md5(string).hexdigest()[:<NUM_LIT:8>], <NUM_LIT:16>) & <NUM_LIT><EOL>

Generate a new seed, from the given seed and salt.


def fill_triangular(x, upper=False, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>m = tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(x.shape, <NUM_LIT:1>)[-<NUM_LIT:1>])<EOL>if m is not None:<EOL><INDENT>m = np.int32(m)<EOL>n = np.sqrt(<NUM_LIT> + <NUM_LIT> * m) - <NUM_LIT:0.5><EOL>if n != np.floor(n):<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>".format(m))<EOL><DEDENT>n = np.int32(n)<EOL>static_final_shape = x.shape[:-<NUM_LIT:1>].concatenate([n, n])<EOL><DEDENT>else:<EOL><INDENT>m = tf.shape(input=x)[-<NUM_LIT:1>]<EOL>n = tf.cast(<EOL>tf.sqrt(<NUM_LIT> + tf.cast(<NUM_LIT:2> * m, dtype=tf.float32)), dtype=tf.int32)<EOL>static_final_shape = tensorshape_util.with_rank_at_least(<EOL>x.shape, <NUM_LIT:1>)[:-<NUM_LIT:1>].concatenate([None, None])<EOL><DEDENT>ndims = prefer_static_rank(x)<EOL>if upper:<EOL><INDENT>x_list = [x, tf.reverse(x[..., n:], axis=[ndims - <NUM_LIT:1>])]<EOL><DEDENT>else:<EOL><INDENT>x_list = [x[..., n:], tf.reverse(x, axis=[ndims - <NUM_LIT:1>])]<EOL><DEDENT>new_shape = (<EOL>tensorshape_util.as_list(static_final_shape)<EOL>if tensorshape_util.is_fully_defined(static_final_shape) else tf.concat(<EOL>[tf.shape(input=x)[:-<NUM_LIT:1>], [n, n]], axis=<NUM_LIT:0>))<EOL>x = tf.reshape(tf.concat(x_list, axis=-<NUM_LIT:1>), new_shape)<EOL>x = tf.linalg.band_part(<EOL>x, num_lower=(<NUM_LIT:0> if upper else -<NUM_LIT:1>), num_upper=(-<NUM_LIT:1> if upper else <NUM_LIT:0>))<EOL>tensorshape_util.set_shape(x, static_final_shape)<EOL>return x<EOL><DEDENT>

r"""Creates a (batch of) triangular matrix from a vector of inputs.

    Created matrix can be lower- or upper-triangular. (It is more efficient to
    create the matrix as upper or lower, rather than transpose.)

    Triangular matrix elements are filled in a clockwise spiral. See example,
    below.

    If `x.shape` is `[b1, b2, ..., bB, d]` then the output shape is
    `[b1, b2, ..., bB, n, n]` where `n` is such that `d = n(n+1)/2`, i.e.,
    `n = int(np.sqrt(0.25 + 2. * m) - 0.5)`.

    Example:

    ```python
    fill_triangular([1, 2, 3, 4, 5, 6])
    # ==> [[4, 0, 0],
    #      [6, 5, 0],
    #      [3, 2, 1]]

    fill_triangular([1, 2, 3, 4, 5, 6], upper=True)
    # ==> [[1, 2, 3],
    #      [0, 5, 6],
    #      [0, 0, 4]]
    ```

    The key trick is to create an upper triangular matrix by concatenating `x`
    and a tail of itself, then reshaping.

    Suppose that we are filling the upper triangle of an `n`-by-`n` matrix `M`
    from a vector `x`. The matrix `M` contains n**2 entries total. The vector `x`
    contains `n * (n+1) / 2` entries. For concreteness, we'll consider `n = 5`
    (so `x` has `15` entries and `M` has `25`). We'll concatenate `x` and `x` with
    the first (`n = 5`) elements removed and reversed:

    ```python
    x = np.arange(15) + 1
    xc = np.concatenate([x, x[5:][::-1]])
    # ==> array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 15, 14, 13,
    #            12, 11, 10, 9, 8, 7, 6])

    # (We add one to the arange result to disambiguate the zeros below the
    # diagonal of our upper-triangular matrix from the first entry in `x`.)

    # Now, when reshapedlay this out as a matrix:
    y = np.reshape(xc, [5, 5])
    # ==> array([[ 1,  2,  3,  4,  5],
    #            [ 6,  7,  8,  9, 10],
    #            [11, 12, 13, 14, 15],
    #            [15, 14, 13, 12, 11],
    #            [10,  9,  8,  7,  6]])

    # Finally, zero the elements below the diagonal:
    y = np.triu(y, k=0)
    # ==> array([[ 1,  2,  3,  4,  5],
    #            [ 0,  7,  8,  9, 10],
    #            [ 0,  0, 13, 14, 15],
    #            [ 0,  0,  0, 12, 11],
    #            [ 0,  0,  0,  0,  6]])
    ```

    From this example we see that the resuting matrix is upper-triangular, and
    contains all the entries of x, as desired. The rest is details:
    - If `n` is even, `x` doesn't exactly fill an even number of rows (it fills
      `n / 2` rows and half of an additional row), but the whole scheme still
      works.
    - If we want a lower triangular matrix instead of an upper triangular,
      we remove the first `n` elements from `x` rather than from the reversed
      `x`.

    For additional comparisons, a pure numpy version of this function can be found
    in `distribution_util_test.py`, function `_fill_triangular`.

    Args:
      x: `Tensor` representing lower (or upper) triangular elements.
      upper: Python `bool` representing whether output matrix should be upper
        triangular (`True`) or lower triangular (`False`, default).
      name: Python `str`. The name to give this op.

    Returns:
      tril: `Tensor` with lower (or upper) triangular elements filled from `x`.

    Raises:
      ValueError: if `x` cannot be mapped to a triangular matrix.


def fill_triangular_inverse(x, upper=False, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>n = tf.compat.dimension_value(<EOL>tensorshape_util.with_rank_at_least(x.shape, <NUM_LIT:2>)[-<NUM_LIT:1>])<EOL>if n is not None:<EOL><INDENT>n = np.int32(n)<EOL>m = np.int32((n * (n + <NUM_LIT:1>)) // <NUM_LIT:2>)<EOL>static_final_shape = x.shape[:-<NUM_LIT:2>].concatenate([m])<EOL><DEDENT>else:<EOL><INDENT>n = tf.shape(input=x)[-<NUM_LIT:1>]<EOL>m = (n * (n + <NUM_LIT:1>)) // <NUM_LIT:2><EOL>static_final_shape = tensorshape_util.with_rank_at_least(<EOL>x.shape, <NUM_LIT:2>)[:-<NUM_LIT:2>].concatenate([None])<EOL><DEDENT>ndims = prefer_static_rank(x)<EOL>if upper:<EOL><INDENT>initial_elements = x[..., <NUM_LIT:0>, :]<EOL>triangular_portion = x[..., <NUM_LIT:1>:, :]<EOL><DEDENT>else:<EOL><INDENT>initial_elements = tf.reverse(x[..., -<NUM_LIT:1>, :], axis=[ndims - <NUM_LIT:2>])<EOL>triangular_portion = x[..., :-<NUM_LIT:1>, :]<EOL><DEDENT>rotated_triangular_portion = tf.reverse(<EOL>tf.reverse(triangular_portion, axis=[ndims - <NUM_LIT:1>]), axis=[ndims - <NUM_LIT:2>])<EOL>consolidated_matrix = triangular_portion + rotated_triangular_portion<EOL>end_sequence = tf.reshape(<EOL>consolidated_matrix,<EOL>tf.concat([tf.shape(input=x)[:-<NUM_LIT:2>], [n * (n - <NUM_LIT:1>)]], axis=<NUM_LIT:0>))<EOL>y = tf.concat([initial_elements, end_sequence[..., :m - n]], axis=-<NUM_LIT:1>)<EOL>tensorshape_util.set_shape(y, static_final_shape)<EOL>return y<EOL><DEDENT>

Creates a vector from a (batch of) triangular matrix.

    The vector is created from the lower-triangular or upper-triangular portion
    depending on the value of the parameter `upper`.

    If `x.shape` is `[b1, b2, ..., bB, n, n]` then the output shape is
    `[b1, b2, ..., bB, d]` where `d = n (n + 1) / 2`.

    Example:

    ```python
    fill_triangular_inverse(
      [[4, 0, 0],
       [6, 5, 0],
       [3, 2, 1]])

    # ==> [1, 2, 3, 4, 5, 6]

    fill_triangular_inverse(
      [[1, 2, 3],
       [0, 5, 6],
       [0, 0, 4]], upper=True)

    # ==> [1, 2, 3, 4, 5, 6]
    ```

    Args:
      x: `Tensor` representing lower (or upper) triangular elements.
      upper: Python `bool` representing whether output matrix should be upper
        triangular (`True`) or lower triangular (`False`, default).
      name: Python `str`. The name to give this op.

    Returns:
      flat_tril: (Batch of) vector-shaped `Tensor` representing vectorized lower
        (or upper) triangular elements from `x`.


def tridiag(below=None, diag=None, above=None, name=None):

def _pad(x):<EOL><INDENT>"""<STR_LIT>"""<EOL>shape = tf.concat([tf.shape(input=x)[:-<NUM_LIT:1>], [<NUM_LIT:1>]], axis=<NUM_LIT:0>)<EOL>z = tf.zeros(shape, dtype=x.dtype)<EOL>return tf.concat([z, x, z], axis=-<NUM_LIT:1>)<EOL><DEDENT>def _add(*x):<EOL><INDENT>"""<STR_LIT>"""<EOL>s = None<EOL>for y in x:<EOL><INDENT>if y is None:<EOL><INDENT>continue<EOL><DEDENT>elif s is None:<EOL><INDENT>s = y<EOL><DEDENT>else:<EOL><INDENT>s += y<EOL><DEDENT><DEDENT>if s is None:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>return s<EOL><DEDENT>with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>if below is not None:<EOL><INDENT>below = tf.convert_to_tensor(value=below, name="<STR_LIT>")<EOL>below = tf.linalg.diag(_pad(below))[..., :-<NUM_LIT:1>, <NUM_LIT:1>:]<EOL><DEDENT>if diag is not None:<EOL><INDENT>diag = tf.convert_to_tensor(value=diag, name="<STR_LIT>")<EOL>diag = tf.linalg.diag(diag)<EOL><DEDENT>if above is not None:<EOL><INDENT>above = tf.convert_to_tensor(value=above, name="<STR_LIT>")<EOL>above = tf.linalg.diag(_pad(above))[..., <NUM_LIT:1>:, :-<NUM_LIT:1>]<EOL><DEDENT>return _add(below, diag, above)<EOL><DEDENT>

Creates a matrix with values set above, below, and on the diagonal.

    Example:

    ```python
    tridiag(below=[1., 2., 3.],
            diag=[4., 5., 6., 7.],
            above=[8., 9., 10.])
    # ==> array([[  4.,   8.,   0.,   0.],
    #            [  1.,   5.,   9.,   0.],
    #            [  0.,   2.,   6.,  10.],
    #            [  0.,   0.,   3.,   7.]], dtype=float32)
    ```

    Warning: This Op is intended for convenience, not efficiency.

    Args:
      below: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the below
        diagonal part. `None` is logically equivalent to `below = 0`.
      diag: `Tensor` of shape `[B1, ..., Bb, d]` corresponding to the diagonal
        part.  `None` is logically equivalent to `diag = 0`.
      above: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the above
        diagonal part.  `None` is logically equivalent to `above = 0`.
      name: Python `str`. The name to give this op.

    Returns:
      tridiag: `Tensor` with values set above, below and on the diagonal.

    Raises:
      ValueError: if all inputs are `None`.