with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[sparse_or_dense_matrix, dense_vector]):<EOL><INDENT>dense_vector = tf.convert_to_tensor(<EOL>value=dense_vector, dtype_hint=tf.float32, name='<STR_LIT>')<EOL>return tf.squeeze(<EOL>sparse_or_dense_matmul(<EOL>sparse_or_dense_matrix,<EOL>dense_vector[..., tf.newaxis],<EOL>validate_args=validate_args,<EOL>**kwargs),<EOL>axis=[-<NUM_LIT:1>])<EOL><DEDENT>

Returns (batched) matmul of a (sparse) matrix with a column vector.

    Args:
      sparse_or_dense_matrix: `SparseTensor` or `Tensor` representing a (batch of)
        matrices.
      dense_vector: `Tensor` representing a (batch of) vectors, with the same
        batch shape as `sparse_or_dense_matrix`. The shape must be compatible with
        the shape of `sparse_or_dense_matrix` and kwargs.
      validate_args: When `True`, additional assertions might be embedded in the
        graph.
        Default value: `False` (i.e., no graph assertions are added).
      name: Python `str` prefixed to ops created by this function.
        Default value: "sparse_or_dense_matvecmul".
      **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul` or
        `tf.matmul`.

    Returns:
      product: A dense (batch of) vector-shaped Tensor of the same batch shape and
      dtype as `sparse_or_dense_matrix` and `dense_vector`.


def _sparse_tensor_dense_matmul(sp_a, b, **kwargs):

batch_shape = _get_shape(sp_a)[:-<NUM_LIT:2>]<EOL>sp_a = tf.sparse.reshape(sp_a, tf.concat([[-<NUM_LIT:1>], _get_shape(sp_a)[-<NUM_LIT:2>:]],<EOL>axis=<NUM_LIT:0>))<EOL>b = tf.reshape(b, tf.concat([[-<NUM_LIT:1>], _get_shape(b)[-<NUM_LIT:1>:]], axis=<NUM_LIT:0>))<EOL>out = tf.sparse.sparse_dense_matmul(_sparse_block_diag(sp_a), b, **kwargs)<EOL>return tf.reshape(<EOL>out,<EOL>tf.concat([batch_shape, [-<NUM_LIT:1>], _get_shape(out)[-<NUM_LIT:1>:]], axis=<NUM_LIT:0>))<EOL>

Returns (batched) matmul of a SparseTensor with a Tensor.

    Args:
      sp_a: `SparseTensor` representing a (batch of) matrices.
      b: `Tensor` representing a (batch of) matrices, with the same batch shape of
        `sp_a`. The shape must be compatible with the shape of `sp_a` and kwargs.
      **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul`.

    Returns:
      product: A dense (batch of) matrix-shaped Tensor of the same batch shape and
      dtype as `sp_a` and `b`. If `sp_a` or `b` is adjointed through `kwargs` then
      the shape is adjusted accordingly.


<EOL>sp_a_shape = tf.convert_to_tensor(value=_get_shape(sp_a, tf.int64))<EOL>ind_mat = tf.concat([[sp_a_shape[-<NUM_LIT:2>:]], tf.eye(<NUM_LIT:2>, dtype=tf.int64)], axis=<NUM_LIT:0>)<EOL>indices = tf.matmul(sp_a.indices, ind_mat)<EOL>dense_shape = sp_a_shape[<NUM_LIT:0>] * sp_a_shape[<NUM_LIT:1>:]<EOL>return tf.SparseTensor(<EOL>indices=indices, values=sp_a.values, dense_shape=dense_shape)<EOL>

Returns a block diagonal rank 2 SparseTensor from a batch of SparseTensors.

    Args:
      sp_a: A rank 3 `SparseTensor` representing a batch of matrices.

    Returns:
      sp_block_diag_a: matrix-shaped, `float` `SparseTensor` with the same dtype
      as `sparse_or_matrix`, of shape [B * M, B * N] where `sp_a` has shape
      [B, M, N]. Each [M, N] batch of `sp_a` is lined up along the diagonal.


assertions = []<EOL>if not a.dtype.is_floating:<EOL><INDENT>raise TypeError('<STR_LIT>'<EOL>'<STR_LIT>'.format(a.dtype.name))<EOL><DEDENT>if a.shape.ndims is not None:<EOL><INDENT>if a.shape.ndims < <NUM_LIT:2>:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'.format(a.shape.ndims))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(tf.compat.v1.assert_rank_at_least(<EOL>a, rank=<NUM_LIT:2>, message='<STR_LIT>'))<EOL><DEDENT>return assertions<EOL>


positive_axis = axis if axis >= <NUM_LIT:0> else axis + len(y_ref_shape)<EOL>x_ref_min = <NUM_LIT:0.><EOL>x_ref_max = <NUM_LIT:2> * np.pi<EOL>implied_x_ref_1d = np.linspace(<EOL>x_ref_min, x_ref_max, y_ref_shape[axis], dtype=dtype)<EOL>implied_x_ref = implied_x_ref_1d.copy()<EOL>for _ in y_ref_shape[positive_axis + <NUM_LIT:1>:]:<EOL><INDENT>implied_x_ref = implied_x_ref[..., np.newaxis]<EOL><DEDENT>for _ in y_ref_shape[:positive_axis]:<EOL><INDENT>implied_x_ref = implied_x_ref[np.newaxis, ...]<EOL><DEDENT>implied_x_ref = implied_x_ref + np.zeros(y_ref_shape, dtype=dtype)<EOL>y_ref = np.sin(implied_x_ref)<EOL>x = np.linspace(<EOL>x_ref_min - <NUM_LIT:0.5>, x_ref_max + <NUM_LIT:0.5>, np.prod(x_shape),<EOL>dtype=dtype).reshape(x_shape)<EOL>y = tfp.math.interp_regular_1d_grid(<EOL>x,<EOL>x_ref_min,<EOL>x_ref_max,<EOL>y_ref,<EOL>fill_value=fill_value,<EOL>axis=axis)<EOL>expected_out_shape = (<EOL>y_ref.shape[:positive_axis] + x_shape +<EOL>y_ref.shape[positive_axis + <NUM_LIT:1>:])<EOL>self.assertAllEqual(expected_out_shape, y.shape)<EOL>sp_func = scipy_interpolate.interp1d(<EOL>implied_x_ref_1d,<EOL>y_ref,<EOL>kind='<STR_LIT>',<EOL>axis=axis,<EOL>bounds_error=False,  <EOL>fill_value=fill_value)<EOL>y_expected = sp_func(x)<EOL>self.assertAllClose(self.evaluate(y), y_expected, atol=<NUM_LIT>, rtol=<NUM_LIT:0>)<EOL>

Check correctness on a sinusoidal function.


def dense_to_sparse(x, ignore_value=None, name=None):

<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>', [x, ignore_value]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>if ignore_value is None:<EOL><INDENT>if x.dtype.base_dtype == tf.string:<EOL><INDENT>ignore_value = '<STR_LIT>'<EOL><DEDENT>else:<EOL><INDENT>ignore_value = x.dtype.as_numpy_dtype(<NUM_LIT:0>)<EOL><DEDENT>ignore_value = tf.cast(ignore_value, x.dtype, name='<STR_LIT>')<EOL><DEDENT>indices = tf.where(tf.not_equal(x, ignore_value), name='<STR_LIT>')<EOL>return tf.SparseTensor(<EOL>indices=indices,<EOL>values=tf.gather_nd(x, indices, name='<STR_LIT>'),<EOL>dense_shape=tf.shape(input=x, out_type=tf.int64, name='<STR_LIT>'))<EOL><DEDENT>

Converts dense `Tensor` to `SparseTensor`, dropping `ignore_value` cells.

    Args:
      x: A `Tensor`.
      ignore_value: Entries in `x` equal to this value will be
        absent from the return `SparseTensor`. If `None`, default value of
        `x` dtype will be used (e.g. '' for `str`, 0 for `int`).
      name: Python `str` prefix for ops created by this function.

    Returns:
      sparse_x: A `tf.SparseTensor` with the same shape as `x`.

    Raises:
      ValueError: when `x`'s rank is `None`.


<EOL>with tf.compat.v1.name_scope(<EOL>name,<EOL>values=[<EOL>x, x_ref_min, x_ref_max, y_ref, axis, fill_value, fill_value_below,<EOL>fill_value_above<EOL>]):<EOL><INDENT>allowed_fv_st = ('<STR_LIT>', '<STR_LIT>')<EOL>for fv in (fill_value, fill_value_below, fill_value_above):<EOL><INDENT>if isinstance(fv, str) and fv not in allowed_fv_st:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(<EOL>fv, allowed_fv_st))<EOL><DEDENT><DEDENT>need_separate_fills = (<EOL>fill_value_above is not None or fill_value_below is not None or<EOL>fill_value == '<STR_LIT>'  <EOL>)<EOL>if need_separate_fills and fill_value_above is None:<EOL><INDENT>fill_value_above = fill_value<EOL><DEDENT>if need_separate_fills and fill_value_below is None:<EOL><INDENT>fill_value_below = fill_value<EOL><DEDENT>dtype = dtype_util.common_dtype([x, x_ref_min, x_ref_max, y_ref],<EOL>preferred_dtype=tf.float32)<EOL>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>', dtype=dtype)<EOL>x_ref_min = tf.convert_to_tensor(<EOL>value=x_ref_min, name='<STR_LIT>', dtype=dtype)<EOL>x_ref_max = tf.convert_to_tensor(<EOL>value=x_ref_max, name='<STR_LIT>', dtype=dtype)<EOL>if not batch_y_ref:<EOL><INDENT>_assert_ndims_statically(x_ref_min, expect_ndims=<NUM_LIT:0>)<EOL>_assert_ndims_statically(x_ref_max, expect_ndims=<NUM_LIT:0>)<EOL><DEDENT>y_ref = tf.convert_to_tensor(value=y_ref, name='<STR_LIT>', dtype=dtype)<EOL>if batch_y_ref:<EOL><INDENT>x_ref_min = x_ref_min[..., tf.newaxis]<EOL>x_ref_max = x_ref_max[..., tf.newaxis]<EOL><DEDENT>axis = tf.convert_to_tensor(value=axis, name='<STR_LIT>', dtype=tf.int32)<EOL>axis = distribution_util.make_non_negative_axis(axis, tf.rank(y_ref))<EOL>_assert_ndims_statically(axis, expect_ndims=<NUM_LIT:0>)<EOL>ny = tf.cast(tf.shape(input=y_ref)[axis], dtype)<EOL>if grid_regularizing_transform is None:<EOL><INDENT>g = lambda x: x<EOL><DEDENT>else:<EOL><INDENT>g = grid_regularizing_transform<EOL><DEDENT>fractional_idx = ((g(x) - g(x_ref_min)) / (g(x_ref_max) - g(x_ref_min)))<EOL>x_idx_unclipped = fractional_idx * (ny - <NUM_LIT:1>)<EOL>nan_idx = tf.math.is_nan(x_idx_unclipped)<EOL>x_idx_unclipped = tf.where(nan_idx, tf.zeros_like(x_idx_unclipped),<EOL>x_idx_unclipped)<EOL>x_idx = tf.clip_by_value(x_idx_unclipped, tf.zeros((), dtype=dtype), ny - <NUM_LIT:1>)<EOL>idx_below = tf.floor(x_idx)<EOL>idx_above = tf.minimum(idx_below + <NUM_LIT:1>, ny - <NUM_LIT:1>)<EOL>idx_below = tf.maximum(idx_above - <NUM_LIT:1>, <NUM_LIT:0>)<EOL>idx_below_int32 = tf.cast(idx_below, dtype=tf.int32)<EOL>idx_above_int32 = tf.cast(idx_above, dtype=tf.int32)<EOL>if batch_y_ref:<EOL><INDENT>y_ref_below = _batch_gather_with_broadcast(y_ref, idx_below_int32, axis)<EOL>y_ref_above = _batch_gather_with_broadcast(y_ref, idx_above_int32, axis)<EOL><DEDENT>else:<EOL><INDENT>y_ref_below = tf.gather(y_ref, idx_below_int32, axis=axis)<EOL>y_ref_above = tf.gather(y_ref, idx_above_int32, axis=axis)<EOL><DEDENT>t = x_idx - idx_below<EOL>if batch_y_ref:<EOL><INDENT>expand_x_fn = _make_expand_x_fn_for_batch_interpolation(y_ref, axis)<EOL><DEDENT>else:<EOL><INDENT>expand_x_fn = _make_expand_x_fn_for_non_batch_interpolation(y_ref, axis)<EOL><DEDENT>t = expand_x_fn(t)<EOL>nan_idx = expand_x_fn(nan_idx, broadcast=True)<EOL>x_idx_unclipped = expand_x_fn(x_idx_unclipped, broadcast=True)<EOL>y = t * y_ref_above + (<NUM_LIT:1> - t) * y_ref_below<EOL>y = tf.where(nan_idx,<EOL>tf.fill(tf.shape(input=y), tf.constant(np.nan, y.dtype)), y)<EOL>if not need_separate_fills:<EOL><INDENT>if fill_value == '<STR_LIT>':<EOL><INDENT>pass  <EOL><DEDENT>else:<EOL><INDENT>y = tf.where((x_idx_unclipped < <NUM_LIT:0>) | (x_idx_unclipped > ny - <NUM_LIT:1>),<EOL>fill_value + tf.zeros_like(y), y)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>if fill_value_below == '<STR_LIT>':<EOL><INDENT>pass  <EOL><DEDENT>elif fill_value_below == '<STR_LIT>':<EOL><INDENT>if batch_y_ref:<EOL><INDENT>y_0 = tf.gather(y_ref, [<NUM_LIT:0>], axis=axis)<EOL>y_1 = tf.gather(y_ref, [<NUM_LIT:1>], axis=axis)<EOL><DEDENT>else:<EOL><INDENT>y_0 = tf.gather(<EOL>y_ref, tf.zeros(tf.shape(input=x), dtype=tf.int32), axis=axis)<EOL>y_1 = tf.gather(<EOL>y_ref, tf.ones(tf.shape(input=x), dtype=tf.int32), axis=axis)<EOL><DEDENT>x_delta = (x_ref_max - x_ref_min) / (ny - <NUM_LIT:1>)<EOL>x_factor = expand_x_fn((x - x_ref_min) / x_delta, broadcast=True)<EOL>y = tf.where(x_idx_unclipped < <NUM_LIT:0>, y_0 + x_factor * (y_1 - y_0), y)<EOL><DEDENT>else:<EOL><INDENT>y = tf.where(x_idx_unclipped < <NUM_LIT:0>, fill_value_below + tf.zeros_like(y),<EOL>y)<EOL><DEDENT>if fill_value_above == '<STR_LIT>':<EOL><INDENT>pass  <EOL><DEDENT>elif fill_value_above == '<STR_LIT>':<EOL><INDENT>ny_int32 = tf.shape(input=y_ref)[axis]<EOL>if batch_y_ref:<EOL><INDENT>y_n1 = tf.gather(y_ref, [tf.shape(input=y_ref)[axis] - <NUM_LIT:1>], axis=axis)<EOL>y_n2 = tf.gather(y_ref, [tf.shape(input=y_ref)[axis] - <NUM_LIT:2>], axis=axis)<EOL><DEDENT>else:<EOL><INDENT>y_n1 = tf.gather(<EOL>y_ref, tf.fill(tf.shape(input=x), ny_int32 - <NUM_LIT:1>), axis=axis)<EOL>y_n2 = tf.gather(<EOL>y_ref, tf.fill(tf.shape(input=x), ny_int32 - <NUM_LIT:2>), axis=axis)<EOL><DEDENT>x_delta = (x_ref_max - x_ref_min) / (ny - <NUM_LIT:1>)<EOL>x_factor = expand_x_fn((x - x_ref_max) / x_delta, broadcast=True)<EOL>y = tf.where(x_idx_unclipped > ny - <NUM_LIT:1>,<EOL>y_n1 + x_factor * (y_n1 - y_n2),<EOL>y)<EOL><DEDENT>else:<EOL><INDENT>y = tf.where(x_idx_unclipped > ny - <NUM_LIT:1>,<EOL>fill_value_above + tf.zeros_like(y), y)<EOL><DEDENT><DEDENT>return y<EOL><DEDENT>

1-D interpolation that works with/without batching.


Linear `1-D` interpolation on a regular (constant spacing) grid.

    Given reference values, this function computes a piecewise linear interpolant
    and evaluates it on a new set of `x` values.

    The interpolant is built from `C` reference values indexed by one dimension
    of `y_ref` (specified by the `axis` kwarg).

    If `y_ref` is a vector, then each value `y_ref[i]` is considered to be equal
    to `f(x_ref[i])`, for `C` (implicitly defined) reference values between
    `x_ref_min` and `x_ref_max`:

    ```none
    x_ref[i] = x_ref_min + i * (x_ref_max - x_ref_min) / (C - 1),
    i = 0, ..., C - 1.
    ```

    If `rank(y_ref) > 1`, then dimension `axis` indexes `C` reference values of
    a shape `y_ref.shape[:axis] + y_ref.shape[axis + 1:]` `Tensor`.

    If `rank(x) > 1`, then the output is obtained by effectively flattening `x`,
    interpolating along `axis`, then expanding the result to shape
    `y_ref.shape[:axis] + x.shape + y_ref.shape[axis + 1:]`.

    These shape semantics are equivalent to `scipy.interpolate.interp1d`.

    Args:
      x: Numeric `Tensor` The x-coordinates of the interpolated output values.
      x_ref_min:  Scalar `Tensor` of same `dtype` as `x`.  The minimum value of
        the (implicitly defined) reference `x_ref`.
      x_ref_max:  Scalar `Tensor` of same `dtype` as `x`.  The maximum value of
        the (implicitly defined) reference `x_ref`.
      y_ref:  `N-D` `Tensor` (`N > 0`) of same `dtype` as `x`. The reference
        output values.
      axis:  Scalar `Tensor` designating the dimension of `y_ref` that indexes
        values of the interpolation table.
        Default value: `-1`, the rightmost axis.
      fill_value:  Determines what values output should take for `x` values that
        are below `x_ref_min` or above `x_ref_max`. `Tensor` or one of the strings
        "constant_extension" ==> Extend as constant function. "extrapolate" ==>
        Extrapolate in a linear fashion.
        Default value: `"constant_extension"`
      fill_value_below:  Optional override of `fill_value` for `x < x_ref_min`.
      fill_value_above:  Optional override of `fill_value` for `x > x_ref_max`.
      grid_regularizing_transform:  Optional transformation `g` which regularizes
        the implied spacing of the x reference points.  In other words, if
        provided, we assume `g(x_ref_i)` is a regular grid between `g(x_ref_min)`
        and `g(x_ref_max)`.
      name:  A name to prepend to created ops.
        Default value: `"interp_regular_1d_grid"`.

    Returns:
      y_interp:  Interpolation between members of `y_ref`, at points `x`.
        `Tensor` of same `dtype` as `x`, and shape
        `y.shape[:axis] + x.shape + y.shape[axis + 1:]`

    Raises:
      ValueError:  If `fill_value` is not an allowed string.
      ValueError:  If `axis` is not a scalar.

    #### Examples

    Interpolate a function of one variable:

    ```python
    y_ref = tf.exp(tf.linspace(start=0., stop=10., num=200))

    interp_regular_1d_grid(
        x=[6.0, 0.5, 3.3], x_ref_min=0., x_ref_max=1., y_ref=y_ref)
    ==> approx [exp(6.0), exp(0.5), exp(3.3)]
    ```

    Interpolate a matrix-valued function of one variable:

    ```python
    mat_0 = [[1., 0.], [0., 1.]]
    mat_1 = [[0., -1], [1, 0]]
    y_ref = [mat_0, mat_1]

    # Get three output matrices at once.
    tfp.math.interp_regular_1d_grid(
        x=[0., 0.5, 1.], x_ref_min=0., x_ref_max=1., y_ref=y_ref, axis=0)
    ==> [mat_0, 0.5 * mat_0 + 0.5 * mat_1, mat_1]
    ```

    Interpolate a scalar valued function, and get a matrix of results:

    ```python
    y_ref = tf.exp(tf.linspace(start=0., stop=10., num=200))
    x = [[1.1, 1.2], [2.1, 2.2]]
    tfp.math.interp_regular_1d_grid(x, x_ref_min=0., x_ref_max=10., y_ref=y_ref)
    ==> tf.exp(x)
    ```

    Interpolate a function of one variable on a log-spaced grid:

    ```python
    x_ref = tf.exp(tf.linspace(tf.log(1.), tf.log(100000.), num_pts))
    y_ref = tf.log(x_ref + x_ref**2)

    interp_regular_1d_grid(x=[1.1, 2.2], x_ref_min=1., x_ref_max=100000., y_ref,
        grid_regularizing_transform=tf.log)
    ==> [tf.log(1.1 + 1.1**2), tf.log(2.2 + 2.2**2)]
    ```


Linear `1-D` interpolation on a regular (constant spacing) grid.

    Given [batch of] reference values, this function computes a piecewise linear
    interpolant and evaluates it on a [batch of] of new `x` values.

    The interpolant is built from `C` reference values indexed by one dimension
    of `y_ref` (specified by the `axis` kwarg).

    If `y_ref` is a vector, then each value `y_ref[i]` is considered to be equal
    to `f(x_ref[i])`, for `C` (implicitly defined) reference values between
    `x_ref_min` and `x_ref_max`:

    ```none
    x_ref[i] = x_ref_min + i * (x_ref_max - x_ref_min) / (C - 1),
    i = 0, ..., C - 1.
    ```

    In the general case, dimensions to the left of `axis` in `y_ref` are broadcast
    with leading dimensions in `x`, `x_ref_min`, `x_ref_max`.

    Args:
      x: Numeric `Tensor` The x-coordinates of the interpolated output values for
        each batch.  Shape broadcasts with `[A1, ..., AN, D]`, `N >= 0`.
      x_ref_min:  `Tensor` of same `dtype` as `x`.  The minimum value of the each
        batch of the (implicitly defined) reference `x_ref`. Shape broadcasts with
        `[A1, ..., AN]`, `N >= 0`.
      x_ref_max:  `Tensor` of same `dtype` as `x`.  The maximum value of the each
        batch of the (implicitly defined) reference `x_ref`. Shape broadcasts with
        `[A1, ..., AN]`, `N >= 0`.
      y_ref:  `Tensor` of same `dtype` as `x`.  The reference output values.
        `y_ref.shape[:axis]` broadcasts with the batch shape `[A1, ..., AN]`, and
        `y_ref.shape[axis:]` is `[C, B1, ..., BM]`, so the trailing dimensions
          index `C` reference values of a rank `M` `Tensor` (`M >= 0`).
      axis:  Scalar `Tensor` designating the dimension of `y_ref` that indexes
        values of the interpolation table.
        Default value: `-1`, the rightmost axis.
      fill_value:  Determines what values output should take for `x` values that
        are below `x_ref_min` or above `x_ref_max`. `Tensor` or one of the strings
        "constant_extension" ==> Extend as constant function. "extrapolate" ==>
        Extrapolate in a linear fashion.
        Default value: `"constant_extension"`
      fill_value_below:  Optional override of `fill_value` for `x < x_ref_min`.
      fill_value_above:  Optional override of `fill_value` for `x > x_ref_max`.
      grid_regularizing_transform:  Optional transformation `g` which regularizes
        the implied spacing of the x reference points.  In other words, if
        provided, we assume `g(x_ref_i)` is a regular grid between `g(x_ref_min)`
        and `g(x_ref_max)`.
      name:  A name to prepend to created ops.
        Default value: `"batch_interp_regular_1d_grid"`.

    Returns:
      y_interp:  Interpolation between members of `y_ref`, at points `x`.
        `Tensor` of same `dtype` as `x`, and shape `[A1, ..., AN, D, B1, ..., BM]`

    Raises:
      ValueError:  If `fill_value` is not an allowed string.
      ValueError:  If `axis` is not a scalar.

    #### Examples

    Interpolate a function of one variable:

    ```python
    y_ref = tf.exp(tf.linspace(start=0., stop=10., 20))

    batch_interp_regular_1d_grid(
        x=[6.0, 0.5, 3.3], x_ref_min=0., x_ref_max=1., y_ref=y_ref)
    ==> approx [exp(6.0), exp(0.5), exp(3.3)]
    ```

    Interpolate a batch of functions of one variable.

    ```python
    # First batch member is an exponential function, second is a log.
    implied_x_ref = [tf.linspace(-3., 3.2, 200), tf.linspace(0.5, 3., 200)]
    y_ref = tf.stack(  # Shape [2, 200], 2 batches, 200 reference values per batch
        [tf.exp(implied_x_ref[0]), tf.log(implied_x_ref[1])], axis=0)

    x = [[-1., 1., 0.],  # Shape [2, 3], 2 batches, 3 values per batch.
         [1., 2., 3.]]

    y = tfp.math.batch_interp_regular_1d_grid(  # Shape [2, 3]
        x,
        x_ref_min=[-3., 0.5],
        x_ref_max=[3.2, 3.],
        y_ref=y_ref,
        axis=-1)

    # y[0] approx tf.exp(x[0])
    # y[1] approx tf.log(x[1])
    ```

    Interpolate a function of one variable on a log-spaced grid:

    ```python
    x_ref = tf.exp(tf.linspace(tf.log(1.), tf.log(100000.), num_pts))
    y_ref = tf.log(x_ref + x_ref**2)

    batch_interp_regular_1d_grid(x=[1.1, 2.2], x_ref_min=1., x_ref_max=100000.,
        y_ref, grid_regularizing_transform=tf.log)
    ==> [tf.log(1.1 + 1.1**2), tf.log(2.2 + 2.2**2)]
    ```


with tf.compat.v1.name_scope(<EOL>name,<EOL>default_name='<STR_LIT>',<EOL>values=[x, x_ref_min, x_ref_max, y_ref, fill_value]):<EOL><INDENT>dtype = dtype_util.common_dtype([x, x_ref_min, x_ref_max, y_ref],<EOL>preferred_dtype=tf.float32)<EOL>if isinstance(fill_value, str):<EOL><INDENT>if fill_value != '<STR_LIT>':<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(<EOL>fill_value, '<STR_LIT>'))<EOL><DEDENT><DEDENT>else:<EOL><INDENT>fill_value = tf.convert_to_tensor(<EOL>value=fill_value, name='<STR_LIT>', dtype=dtype)<EOL>_assert_ndims_statically(fill_value, expect_ndims=<NUM_LIT:0>)<EOL><DEDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>', dtype=dtype)<EOL>_assert_ndims_statically(x, expect_ndims_at_least=<NUM_LIT:2>)<EOL>y_ref = tf.convert_to_tensor(value=y_ref, name='<STR_LIT>', dtype=dtype)<EOL>x_ref_min = tf.convert_to_tensor(<EOL>value=x_ref_min, name='<STR_LIT>', dtype=dtype)<EOL>x_ref_max = tf.convert_to_tensor(<EOL>value=x_ref_max, name='<STR_LIT>', dtype=dtype)<EOL>_assert_ndims_statically(<EOL>x_ref_min, expect_ndims_at_least=<NUM_LIT:1>, expect_static=True)<EOL>_assert_ndims_statically(<EOL>x_ref_max, expect_ndims_at_least=<NUM_LIT:1>, expect_static=True)<EOL>nd = tf.compat.dimension_value(x_ref_min.shape[-<NUM_LIT:1>])<EOL>if nd is None:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>x_ref_max.shape[-<NUM_LIT:1>:].assert_is_compatible_with(x_ref_min.shape[-<NUM_LIT:1>:])<EOL>axis = tf.convert_to_tensor(value=axis, dtype=tf.int32, name='<STR_LIT>')<EOL>axis = distribution_util.make_non_negative_axis(axis, tf.rank(y_ref))<EOL>axis.shape.assert_has_rank(<NUM_LIT:0>)<EOL>axis_ = tf.get_static_value(axis)<EOL>y_ref_rank_ = tf.get_static_value(tf.rank(y_ref))<EOL>if axis_ is not None and y_ref_rank_ is not None:<EOL><INDENT>if axis_ + nd > y_ref_rank_:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>axis_, y_ref_rank_, nd))<EOL><DEDENT><DEDENT>x_batch_shape = tf.shape(input=x)[:-<NUM_LIT:2>]<EOL>x_ref_min_batch_shape = tf.shape(input=x_ref_min)[:-<NUM_LIT:1>]<EOL>x_ref_max_batch_shape = tf.shape(input=x_ref_max)[:-<NUM_LIT:1>]<EOL>y_ref_batch_shape = tf.shape(input=y_ref)[:axis]<EOL>batch_shape = y_ref_batch_shape<EOL>for tensor in [x_batch_shape, x_ref_min_batch_shape, x_ref_max_batch_shape]:<EOL><INDENT>batch_shape = tf.broadcast_dynamic_shape(batch_shape, tensor)<EOL><DEDENT>def _batch_of_zeros_with_rightmost_singletons(n_singletons):<EOL><INDENT>"""<STR_LIT>"""<EOL>ones = tf.ones(shape=[n_singletons], dtype=tf.int32)<EOL>return tf.zeros(shape=tf.concat([batch_shape, ones], axis=<NUM_LIT:0>), dtype=dtype)<EOL><DEDENT>x += _batch_of_zeros_with_rightmost_singletons(n_singletons=<NUM_LIT:2>)<EOL>x_ref_min += _batch_of_zeros_with_rightmost_singletons(n_singletons=<NUM_LIT:1>)<EOL>x_ref_max += _batch_of_zeros_with_rightmost_singletons(n_singletons=<NUM_LIT:1>)<EOL>y_ref += _batch_of_zeros_with_rightmost_singletons(<EOL>n_singletons=tf.rank(y_ref) - axis)<EOL>return _batch_interp_with_gather_nd(<EOL>x=x,<EOL>x_ref_min=x_ref_min,<EOL>x_ref_max=x_ref_max,<EOL>y_ref=y_ref,<EOL>nd=nd,<EOL>fill_value=fill_value,<EOL>batch_dims=tf.get_static_value(tf.rank(x)) - <NUM_LIT:2>)<EOL><DEDENT>

Multi-linear interpolation on a regular (constant spacing) grid.

    Given [a batch of] reference values, this function computes a multi-linear
    interpolant and evaluates it on [a batch of] of new `x` values.

    The interpolant is built from reference values indexed by `nd` dimensions
    of `y_ref`, starting at `axis`.

    For example, take the case of a `2-D` scalar valued function and no leading
    batch dimensions.  In this case, `y_ref.shape = [C1, C2]` and `y_ref[i, j]`
    is the reference value corresponding to grid point

    ```
    [x_ref_min[0] + i * (x_ref_max[0] - x_ref_min[0]) / (C1 - 1),
     x_ref_min[1] + j * (x_ref_max[1] - x_ref_min[1]) / (C2 - 1)]
    ```

    In the general case, dimensions to the left of `axis` in `y_ref` are broadcast
    with leading dimensions in `x`, `x_ref_min`, `x_ref_max`.

    Args:
      x: Numeric `Tensor` The x-coordinates of the interpolated output values for
        each batch.  Shape `[..., D, nd]`, designating [a batch of] `D`
        coordinates in `nd` space.  `D` must be `>= 1` and is not a batch dim.
      x_ref_min:  `Tensor` of same `dtype` as `x`.  The minimum values of the
        (implicitly defined) reference `x_ref`.  Shape `[..., nd]`.
      x_ref_max:  `Tensor` of same `dtype` as `x`.  The maximum values of the
        (implicitly defined) reference `x_ref`.  Shape `[..., nd]`.
      y_ref:  `Tensor` of same `dtype` as `x`.  The reference output values. Shape
        `[..., C1, ..., Cnd, B1,...,BM]`, designating [a batch of] reference
        values indexed by `nd` dimensions, of a shape `[B1,...,BM]` valued
        function (for `M >= 0`).
      axis:  Scalar integer `Tensor`.  Dimensions `[axis, axis + nd)` of `y_ref`
        index the interpolation table.  E.g. `3-D` interpolation of a scalar
        valued function requires `axis=-3` and a `3-D` matrix valued function
        requires `axis=-5`.
      fill_value:  Determines what values output should take for `x` values that
        are below `x_ref_min` or above `x_ref_max`. Scalar `Tensor` or
        "constant_extension" ==> Extend as constant function.
        Default value: `"constant_extension"`
      name:  A name to prepend to created ops.
        Default value: `"batch_interp_regular_nd_grid"`.

    Returns:
      y_interp:  Interpolation between members of `y_ref`, at points `x`.
        `Tensor` of same `dtype` as `x`, and shape `[..., D, B1, ..., BM].`

    Raises:
      ValueError:  If `rank(x) < 2` is determined statically.
      ValueError:  If `axis` is not a scalar is determined statically.
      ValueError:  If `axis + nd > rank(y_ref)` is determined statically.

    #### Examples

    Interpolate a function of one variable.

    ```python
    y_ref = tf.exp(tf.linspace(start=0., stop=10., 20))

    tfp.math.batch_interp_regular_nd_grid(
        # x.shape = [3, 1], x_ref_min/max.shape = [1].  Trailing `1` for `1-D`.
        x=[[6.0], [0.5], [3.3]], x_ref_min=[0.], x_ref_max=[1.], y_ref=y_ref)
    ==> approx [exp(6.0), exp(0.5), exp(3.3)]
    ```

    Interpolate a scalar function of two variables.

    ```python
    x_ref_min = [0., 2 * np.pi]
    x_ref_max = [0., 2 * np.pi]

    # Build y_ref.
    x0s, x1s = tf.meshgrid(
        tf.linspace(x_ref_min[0], x_ref_max[0], num=100),
        tf.linspace(x_ref_min[1], x_ref_max[1], num=100),
        indexing='ij')

    def func(x0, x1):
      return tf.sin(x0) * tf.cos(x1)

    y_ref = func(x0s, x1s)

    x = np.pi * tf.random_uniform(shape=(10, 2))

    tfp.math.batch_interp_regular_nd_grid(x, x_ref_min, x_ref_max, y_ref, axis=-2)
    ==> tf.sin(x[:, 0]) * tf.cos(x[:, 1])
    ```


def _batch_interp_with_gather_nd(x, x_ref_min, x_ref_max, y_ref, nd, fill_value,<EOL>batch_dims):

dtype = x.dtype<EOL>ny = tf.cast(tf.shape(input=y_ref)[batch_dims:batch_dims + nd], dtype)<EOL>x_ref_min_expanded = tf.expand_dims(x_ref_min, axis=-<NUM_LIT:2>)<EOL>x_ref_max_expanded = tf.expand_dims(x_ref_max, axis=-<NUM_LIT:2>)<EOL>x_idx_unclipped = (ny - <NUM_LIT:1>) * (x - x_ref_min_expanded) / (<EOL>x_ref_max_expanded - x_ref_min_expanded)<EOL>nan_idx = tf.math.is_nan(x_idx_unclipped)<EOL>x_idx_unclipped = tf.where(nan_idx, tf.zeros_like(x_idx_unclipped),<EOL>x_idx_unclipped)<EOL>x_idx = tf.clip_by_value(x_idx_unclipped, tf.zeros((), dtype=dtype), ny - <NUM_LIT:1>)<EOL>idx_below = tf.floor(x_idx)<EOL>idx_above = tf.minimum(idx_below + <NUM_LIT:1>, ny - <NUM_LIT:1>)<EOL>idx_below = tf.maximum(idx_above - <NUM_LIT:1>, <NUM_LIT:0>)<EOL>idx_below_int32 = tf.cast(idx_below, dtype=tf.int32)<EOL>idx_above_int32 = tf.cast(idx_above, dtype=tf.int32)<EOL>idx_below_list = tf.unstack(idx_below_int32, axis=-<NUM_LIT:1>)<EOL>idx_above_list = tf.unstack(idx_above_int32, axis=-<NUM_LIT:1>)<EOL>t = x_idx - idx_below<EOL>def _expand_x_fn(tensor):<EOL><INDENT>extended_shape = tf.concat([<EOL>tf.shape(input=tensor),<EOL>tf.ones_like(tf.shape(input=y_ref)[batch_dims + nd:])<EOL>],<EOL>axis=<NUM_LIT:0>)<EOL>return tf.reshape(tensor, extended_shape)<EOL><DEDENT>t = _expand_x_fn(t)<EOL>s = <NUM_LIT:1> - t<EOL>nan_idx = _expand_x_fn(nan_idx)<EOL>t = tf.where(nan_idx,<EOL>tf.fill(tf.shape(input=t), tf.constant(np.nan, dtype)),<EOL>t)<EOL>terms = []<EOL>for zero_ones_list in _binary_count(nd):<EOL><INDENT>gather_from_y_ref_idx = []<EOL>opposite_volume_t_idx = []<EOL>opposite_volume_s_idx = []<EOL>for k, zero_or_one in enumerate(zero_ones_list):<EOL><INDENT>if zero_or_one == <NUM_LIT:0>:<EOL><INDENT>gather_from_y_ref_idx.append(idx_below_list[k])<EOL>opposite_volume_s_idx.append(k)<EOL><DEDENT>else:<EOL><INDENT>gather_from_y_ref_idx.append(idx_above_list[k])<EOL>opposite_volume_t_idx.append(k)<EOL><DEDENT><DEDENT>ov_axis = tf.rank(x) - <NUM_LIT:1><EOL>opposite_volume = (<EOL>tf.reduce_prod(<EOL>input_tensor=tf.gather(<EOL>t, indices=tf.cast(opposite_volume_t_idx, dtype=tf.int32),<EOL>axis=ov_axis),<EOL>axis=ov_axis) *<EOL>tf.reduce_prod(<EOL>input_tensor=tf.gather(<EOL>s, indices=tf.cast(opposite_volume_s_idx, dtype=tf.int32),<EOL>axis=ov_axis),<EOL>axis=ov_axis)<EOL>)  <EOL>y_ref_pt = tf.gather_nd(<EOL>y_ref,<EOL>tf.stack(gather_from_y_ref_idx, axis=-<NUM_LIT:1>),<EOL>batch_dims=batch_dims)<EOL>terms.append(y_ref_pt * opposite_volume)<EOL><DEDENT>y = tf.math.add_n(terms)<EOL>if tf.debugging.is_numeric_tensor(fill_value):<EOL><INDENT>oob_idx = tf.reduce_any(<EOL>input_tensor=(x_idx_unclipped < <NUM_LIT:0>) | (x_idx_unclipped > ny - <NUM_LIT:1>),<EOL>axis=-<NUM_LIT:1>)<EOL>oob_idx = _expand_x_fn(oob_idx)  <EOL>oob_idx |= tf.fill(tf.shape(input=y), False)<EOL>y = tf.where(oob_idx, tf.fill(tf.shape(input=y), fill_value), y)<EOL><DEDENT>return y<EOL>

N-D interpolation that works with leading batch dims.


ndims = x.shape.ndims<EOL>if ndims is None:<EOL><INDENT>if expect_static:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(x))<EOL><DEDENT>return<EOL><DEDENT>if expect_ndims is not None and ndims != expect_ndims:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(expect_ndims, ndims))<EOL><DEDENT>if expect_ndims_at_least is not None and ndims < expect_ndims_at_least:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(<EOL>expect_ndims_at_least, ndims))<EOL><DEDENT>

Assert that Tensor x has expected number of dimensions.


<EOL>y_ref_shape = tf.shape(input=y_ref)<EOL>y_ref_shape_left = y_ref_shape[:axis]<EOL>y_ref_shape_right = y_ref_shape[axis + <NUM_LIT:1>:]<EOL>def expand_ends(x, broadcast=False):<EOL><INDENT>"""<STR_LIT>"""<EOL>expanded_shape = tf.pad(<EOL>tensor=tf.shape(input=x),<EOL>paddings=[[axis, tf.size(input=y_ref_shape_right)]],<EOL>constant_values=<NUM_LIT:1>)<EOL>x_expanded = tf.reshape(x, expanded_shape)<EOL>if broadcast:<EOL><INDENT>out_shape = tf.concat((<EOL>y_ref_shape_left,<EOL>tf.shape(input=x),<EOL>y_ref_shape_right,<EOL>),<EOL>axis=<NUM_LIT:0>)<EOL>if x.dtype.is_bool:<EOL><INDENT>x_expanded = x_expanded | tf.cast(tf.zeros(out_shape), tf.bool)<EOL><DEDENT>else:<EOL><INDENT>x_expanded += tf.zeros(out_shape, dtype=x.dtype)<EOL><DEDENT><DEDENT>return x_expanded<EOL><DEDENT>return expand_ends<EOL>

Make func to expand left/right (of axis) dims of tensors shaped like x.


<EOL>y_ref_shape = tf.shape(input=y_ref)<EOL>y_ref_shape_left = y_ref_shape[:axis]<EOL>y_ref_shape_right = y_ref_shape[axis + <NUM_LIT:1>:]<EOL>def expand_right_dims(x, broadcast=False):<EOL><INDENT>"""<STR_LIT>"""<EOL>expanded_shape_left = tf.broadcast_dynamic_shape(<EOL>tf.shape(input=x)[:-<NUM_LIT:1>],<EOL>tf.ones([tf.size(input=y_ref_shape_left)], dtype=tf.int32))<EOL>expanded_shape = tf.concat(<EOL>(expanded_shape_left, tf.shape(input=x)[-<NUM_LIT:1>:],<EOL>tf.ones([tf.size(input=y_ref_shape_right)], dtype=tf.int32)),<EOL>axis=<NUM_LIT:0>)<EOL>x_expanded = tf.reshape(x, expanded_shape)<EOL>if broadcast:<EOL><INDENT>broadcast_shape_left = tf.broadcast_dynamic_shape(<EOL>tf.shape(input=x)[:-<NUM_LIT:1>], y_ref_shape_left)<EOL>broadcast_shape = tf.concat(<EOL>(broadcast_shape_left, tf.shape(input=x)[-<NUM_LIT:1>:], y_ref_shape_right),<EOL>axis=<NUM_LIT:0>)<EOL>if x.dtype.is_bool:<EOL><INDENT>x_expanded = x_expanded | tf.cast(tf.zeros(broadcast_shape), tf.bool)<EOL><DEDENT>else:<EOL><INDENT>x_expanded += tf.zeros(broadcast_shape, dtype=x.dtype)<EOL><DEDENT><DEDENT>return x_expanded<EOL><DEDENT>return expand_right_dims<EOL>


def _batch_gather_with_broadcast(params, indices, axis):

<EOL>leading_bcast_shape = tf.broadcast_dynamic_shape(<EOL>tf.shape(input=params)[:axis],<EOL>tf.shape(input=indices)[:-<NUM_LIT:1>])<EOL>params += tf.zeros(<EOL>tf.concat((leading_bcast_shape, tf.shape(input=params)[axis:]), axis=<NUM_LIT:0>),<EOL>dtype=params.dtype)<EOL>indices += tf.zeros(<EOL>tf.concat((leading_bcast_shape, tf.shape(input=indices)[-<NUM_LIT:1>:]), axis=<NUM_LIT:0>),<EOL>dtype=indices.dtype)<EOL>return tf.compat.v1.batch_gather(params, indices)<EOL>

Like batch_gather, but broadcasts to the left of axis.


return list(itertools.product([<NUM_LIT:0>, <NUM_LIT:1>], repeat=n))<EOL>

Count `n` binary digits from [0...0] to [1...1].


if not callable(stopping_policy_fn):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>position = tf.convert_to_tensor(<EOL>value=initial_position,<EOL>name='<STR_LIT>',<EOL>)<EOL>value_at_position = tf.convert_to_tensor(<EOL>value=value_at_position or objective_fn(position),<EOL>name='<STR_LIT>',<EOL>dtype=position.dtype.base_dtype)<EOL>zero = tf.zeros_like(position)<EOL>position_tolerance = tf.convert_to_tensor(<EOL>value=position_tolerance, name='<STR_LIT>', dtype=position.dtype)<EOL>value_tolerance = tf.convert_to_tensor(<EOL>value=value_tolerance, name='<STR_LIT>', dtype=position.dtype)<EOL>num_iterations = tf.zeros_like(position, dtype=tf.int32)<EOL>max_iterations = tf.convert_to_tensor(value=max_iterations, dtype=tf.int32)<EOL>max_iterations = tf.broadcast_to(<EOL>max_iterations, name='<STR_LIT>', shape=position.shape)<EOL>if next_position is None:<EOL><INDENT>epsilon = tf.constant(<NUM_LIT>, dtype=position.dtype, shape=position.shape)<EOL>step = position * epsilon + tf.sign(position) * epsilon<EOL><DEDENT>else:<EOL><INDENT>step = next_position - initial_position<EOL><DEDENT>finished = tf.constant(False, shape=position.shape)<EOL>def _should_continue(position, value_at_position, num_iterations, step,<EOL>finished):<EOL><INDENT>"""<STR_LIT>"""<EOL>del position, value_at_position, num_iterations, step  <EOL>return ~tf.convert_to_tensor(<EOL>value=stopping_policy_fn(finished), name='<STR_LIT>', dtype=tf.bool)<EOL><DEDENT>def _body(position, value_at_position, num_iterations, step, finished):<EOL><INDENT>"""<STR_LIT>"""<EOL>was_finished = finished | (num_iterations >= max_iterations) | (<EOL>tf.abs(step) < position_tolerance) | (<EOL>tf.abs(value_at_position) < value_tolerance)<EOL>next_position = tf.where(was_finished, position, position + step)<EOL>value_at_next_position = tf.where(was_finished, value_at_position,<EOL>objective_fn(next_position))<EOL>is_finished = tf.equal(value_at_position, value_at_next_position)<EOL>next_position = tf.where(is_finished & ~was_finished,<EOL>(position + next_position) * <NUM_LIT:0.5>, next_position)<EOL>num_iterations = tf.where(is_finished, num_iterations, num_iterations + <NUM_LIT:1>)<EOL>next_step = tf.where(<EOL>is_finished, zero, step * value_at_next_position /<EOL>(value_at_position - value_at_next_position))<EOL>return (next_position, value_at_next_position, num_iterations, next_step,<EOL>is_finished)<EOL><DEDENT>with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>[position, next_position, value_at_position, max_iterations]):<EOL><INDENT>assertions = []<EOL>if validate_args:<EOL><INDENT>assertions += [<EOL>tf.Assert(<EOL>tf.reduce_all(input_tensor=position_tolerance > zero),<EOL>[position_tolerance]),<EOL>tf.Assert(<EOL>tf.reduce_all(input_tensor=value_tolerance > zero),<EOL>[value_tolerance]),<EOL>tf.Assert(<EOL>tf.reduce_all(input_tensor=max_iterations >= num_iterations),<EOL>[max_iterations]),<EOL>]<EOL><DEDENT>with tf.control_dependencies(assertions):<EOL><INDENT>root, value_at_root, num_iterations, _, _ = tf.while_loop(<EOL>cond=_should_continue,<EOL>body=_body,<EOL>loop_vars=[<EOL>position, value_at_position, num_iterations, step, finished<EOL>])<EOL><DEDENT><DEDENT>return RootSearchResults(<EOL>estimated_root=root,<EOL>objective_at_estimated_root=value_at_root,<EOL>num_iterations=num_iterations)<EOL>

r"""Finds root(s) of a function of single variable using the secant method.

    The [secant method](https://en.wikipedia.org/wiki/Secant_method) is a
    root-finding algorithm that uses a succession of roots of secant lines to
    better approximate a root of a function. The secant method can be thought of
    as a finite-difference approximation of Newton's method.

    Args:
      objective_fn: Python callable for which roots are searched. It must be a
        callable of a single variable. `objective_fn` must return a `Tensor` of
        the same shape and dtype as `initial_position`.
      initial_position: `Tensor` or Python float representing the starting
        position. The function will search for roots in the neighborhood of each
        point. The shape of `initial_position` should match that of the input to
        `objective_fn`.
      next_position: Optional `Tensor` representing the next position in the
        search. If specified, this argument must broadcast with the shape of
        `initial_position` and have the same dtype. It will be used to compute the
        first step to take when searching for roots. If not specified, a default
        value will be used instead.
        Default value: `initial_position * (1 + 1e-4) + sign(initial_position) *
          1e-4`.
      value_at_position: Optional `Tensor` or Pyhon float representing the value
        of `objective_fn` at `initial_position`. If specified, this argument must
        have the same shape and dtype as `initial_position`. If not specified, the
        value will be evaluated during the search.
        Default value: None.
      position_tolerance: Optional `Tensor` representing the tolerance for the
        estimated roots. If specified, this argument must broadcast with the shape
        of `initial_position` and have the same dtype.
        Default value: `1e-8`.
      value_tolerance: Optional `Tensor` representing the tolerance used to check
        for roots. If the absolute value of `objective_fn` is smaller than
        `value_tolerance` at a given position, then that position is considered a
        root for the function. If specified, this argument must broadcast with the
        shape of `initial_position` and have the same dtype.
        Default value: `1e-8`.
      max_iterations: Optional `Tensor` or Python integer specifying the maximum
        number of steps to perform for each initial position. Must broadcast with
        the shape of `initial_position`.
        Default value: `50`.
      stopping_policy_fn: Python `callable` controlling the algorithm termination.
        It must be a callable accepting a `Tensor` of booleans with the shape of
        `initial_position` (each denoting whether the search is finished for each
        starting point), and returning a scalar boolean `Tensor` (indicating
        whether the overall search should stop). Typical values are
        `tf.reduce_all` (which returns only when the search is finished for all
        points), and `tf.reduce_any` (which returns as soon as the search is
        finished for any point).
        Default value: `tf.reduce_all` (returns only when the search is finished
          for all points).
      validate_args: Python `bool` indicating whether to validate arguments such
        as `position_tolerance`, `value_tolerance`, and `max_iterations`.
        Default value: `False`.
      name: Python `str` name prefixed to ops created by this function.

    Returns:
      root_search_results: A Python `namedtuple` containing the following items:
        estimated_root: `Tensor` containing the last position explored. If the
          search was successful within the specified tolerance, this position is
          a root of the objective function.
        objective_at_estimated_root: `Tensor` containing the value of the
          objective function at `position`. If the search was successful within
          the specified tolerance, then this is close to 0.
        num_iterations: The number of iterations performed.

    Raises:
      ValueError: if a non-callable `stopping_policy_fn` is passed.

    #### Examples

    ```python
    import tensorflow as tf
    import tensorflow_probability as tfp
    tf.enable_eager_execution()

    # Example 1: Roots of a single function from two different starting points.

    f = lambda x: (63 * x**5 - 70 * x**3 + 15 * x) / 8.
    x = tf.constant([-1, 10], dtype=tf.float64)

    tfp.math.secant_root(objective_fn=f, initial_position=x))
    # ==> RootSearchResults(
        estimated_root=array([-0.90617985, 0.90617985]),
        objective_at_estimated_root=array([-4.81727769e-10, 7.44957651e-10]),
        num_iterations=array([ 7, 24], dtype=int32))

    tfp.math.secant_root(objective_fn=f,
                         initial_position=x,
                         stopping_policy_fn=tf.reduce_any)
    # ==> RootSearchResults(
        estimated_root=array([-0.90617985, 3.27379206]),
        objective_at_estimated_root=array([-4.81727769e-10, 2.66058312e+03]),
        num_iterations=array([7, 8], dtype=int32))

    # Example 2: Roots of a multiplex function from a single starting point.

    def f(x):
      return tf.constant([0., 63. / 8], dtype=tf.float64) * x**5 \
          + tf.constant([5. / 2, -70. / 8], dtype=tf.float64) * x**3 \
          + tf.constant([-3. / 2, 15. / 8], dtype=tf.float64) * x

    x = tf.constant([-1, -1], dtype=tf.float64)

    tfp.math.secant_root(objective_fn=f, initial_position=x)
    # ==> RootSearchResults(
        estimated_root=array([-0.77459667, -0.90617985]),
        objective_at_estimated_root=array([-7.81339438e-11, -4.81727769e-10]),
        num_iterations=array([7, 7], dtype=int32))

    # Example 3: Roots of a multiplex function from two starting points.

    def f(x):
      return tf.constant([0., 63. / 8], dtype=tf.float64) * x**5 \
          + tf.constant([5. / 2, -70. / 8], dtype=tf.float64) * x**3 \
          + tf.constant([-3. / 2, 15. / 8], dtype=tf.float64) * x

    x = tf.constant([[-1, -1], [10, 10]], dtype=tf.float64)

    tfp.math.secant_root(objective_fn=f, initial_position=x)
    # ==> RootSearchResults(
        estimated_root=array([
            [-0.77459667, -0.90617985],
            [ 0.77459667, 0.90617985]]),
        objective_at_estimated_root=array([
            [-7.81339438e-11, -4.81727769e-10],
            [6.66025013e-11, 7.44957651e-10]]),
        num_iterations=array([
            [7, 7],
            [16, 24]], dtype=int32))
    ```


def random_rademacher(shape, dtype=tf.float32, seed=None, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [shape, seed]):<EOL><INDENT>generation_dtype = tf.int64 if tf.as_dtype(dtype) != tf.int32 else tf.int32<EOL>random_bernoulli = tf.random.uniform(<EOL>shape, minval=<NUM_LIT:0>, maxval=<NUM_LIT:2>, dtype=generation_dtype, seed=seed)<EOL>return tf.cast(<NUM_LIT:2> * random_bernoulli - <NUM_LIT:1>, dtype)<EOL><DEDENT>

Generates `Tensor` consisting of `-1` or `+1`, chosen uniformly at random.

    For more details, see [Rademacher distribution](
    https://en.wikipedia.org/wiki/Rademacher_distribution).

    Args:
      shape: Vector-shaped, `int` `Tensor` representing shape of output.
      dtype: (Optional) TF `dtype` representing `dtype` of output.
      seed: (Optional) Python integer to seed the random number generator.
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'random_rademacher').

    Returns:
      rademacher: `Tensor` with specified `shape` and `dtype` consisting of `-1`
        or `+1` chosen uniformly-at-random.


def random_rayleigh(shape, scale=None, dtype=tf.float32, seed=None, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [shape, scale, seed]):<EOL><INDENT>if scale is not None:<EOL><INDENT>scale = tf.convert_to_tensor(value=scale, dtype=dtype, name='<STR_LIT>')<EOL>shape = tf.broadcast_dynamic_shape(shape, tf.shape(input=scale))<EOL><DEDENT>x = tf.sqrt(-<NUM_LIT> * tf.math.log(<EOL>tf.random.uniform(shape, minval=<NUM_LIT:0>, maxval=<NUM_LIT:1>, dtype=dtype, seed=seed)))<EOL>if scale is None:<EOL><INDENT>return x<EOL><DEDENT>return x * scale<EOL><DEDENT>

Generates `Tensor` of positive reals drawn from a Rayleigh distributions.

    The probability density function of a Rayleigh distribution with `scale`
    parameter is given by:

    ```none
    f(x) = x scale**-2 exp(-x**2 0.5 scale**-2)
    ```

    For more details, see [Rayleigh distribution](
    https://en.wikipedia.org/wiki/Rayleigh_distribution)

    Args:
      shape: Vector-shaped, `int` `Tensor` representing shape of output.
      scale: (Optional) Positive `float` `Tensor` representing `Rayleigh` scale.
        Default value: `None` (i.e., `scale = 1.`).
      dtype: (Optional) TF `dtype` representing `dtype` of output.
        Default value: `tf.float32`.
      seed: (Optional) Python integer to seed the random number generator.
        Default value: `None` (i.e., no seed).
      name: Python `str` name prefixed to Ops created by this function.
        Default value: `None` (i.e., 'random_rayleigh').

    Returns:
      rayleigh: `Tensor` with specified `shape` and `dtype` consisting of positive
        real values drawn from a Rayleigh distribution with specified `scale`.


def custom_gradient(fx, gx, x, fx_gx_manually_stopped=False, name=None):

def maybe_stop(x):<EOL><INDENT>if fx_gx_manually_stopped:<EOL><INDENT>return x<EOL><DEDENT>return tf.stop_gradient(x)<EOL><DEDENT>with tf.compat.v1.name_scope(name, '<STR_LIT>', [fx, gx, x]):<EOL><INDENT>fx = tf.convert_to_tensor(value=fx, name='<STR_LIT>')<EOL>with tf.control_dependencies([fx]):<EOL><INDENT>if is_list_like(x):<EOL><INDENT>x = [identity(x_, name='<STR_LIT:x>') for x_ in x]<EOL><DEDENT>else:<EOL><INDENT>x = [identity(x, name='<STR_LIT:x>')]<EOL><DEDENT>if is_list_like(gx):<EOL><INDENT>gx = [identity(gx_, dtype=fx.dtype, name='<STR_LIT>')<EOL>for gx_ in gx]<EOL><DEDENT>else:<EOL><INDENT>gx = [identity(gx, dtype=fx.dtype, name='<STR_LIT>')]<EOL><DEDENT>override_grad = []<EOL>for x_, gx_ in zip(x, gx):<EOL><INDENT>equal_shape = tf.compat.v1.assert_equal(<EOL>tf.shape(input=x_),<EOL>tf.shape(input=gx_),<EOL>message='<STR_LIT>')<EOL>with tf.control_dependencies([equal_shape]):<EOL><INDENT>zeros_like_x_ = x_ - tf.stop_gradient(x_)<EOL>override_grad.append(<EOL>tf.reduce_sum(input_tensor=maybe_stop(gx_) * zeros_like_x_))<EOL><DEDENT><DEDENT>override_grad = sum(override_grad)<EOL>override_grad /= tf.cast(tf.size(input=fx), dtype=fx.dtype.base_dtype)<EOL>return maybe_stop(fx) + override_grad<EOL><DEDENT><DEDENT>

Embeds a custom gradient into a `Tensor`.

    This function works by clever application of `stop_gradient`. I.e., observe
    that:

    ```none
    h(x) = stop_gradient(f(x)) + stop_gradient(g(x)) * (x - stop_gradient(x))
    ```

    is such that `h(x) == stop_gradient(f(x))` and
    `grad[h(x), x] == stop_gradient(g(x)).`

    In addition to scalar-domain/scalar-range functions, this function also
    supports tensor-domain/scalar-range functions.

    Partial Custom Gradient:

    Suppose `h(x) = htilde(x, y)`. Note that `dh/dx = stop(g(x))` but `dh/dy =
    None`. This is because a `Tensor` cannot have only a portion of its gradient
    stopped. To circumvent this issue, one must manually `stop_gradient` the
    relevant portions of `f`, `g`. For example see the unit-test,
    `test_works_correctly_fx_gx_manually_stopped`.

    Args:
      fx: `Tensor`. Output of function evaluated at `x`.
      gx: `Tensor` or list of `Tensor`s. Gradient of function at (each) `x`.
      x: `Tensor` or list of `Tensor`s. Args of evaluation for `f`.
      fx_gx_manually_stopped: Python `bool` indicating that `fx`, `gx` manually
        have `stop_gradient` applied.
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      fx: Floating-type `Tensor` equal to `f(x)` but which has gradient
        `stop_gradient(g(x))`.


if a.ndim < <NUM_LIT:3>:<EOL><INDENT>return np.linalg.pinv(a, rcond)<EOL><DEDENT>if rcond is None:<EOL><INDENT>rcond = <NUM_LIT> * max(a.shape[-<NUM_LIT:2>], a.shape[-<NUM_LIT:1>]) * np.finfo(a.dtype).eps<EOL><DEDENT>s = np.concatenate([a.shape[:-<NUM_LIT:2>], [a.shape[-<NUM_LIT:1>], a.shape[-<NUM_LIT:2>]]])<EOL>a_pinv = np.zeros(s, dtype=a.dtype)<EOL>for i in np.ndindex(a.shape[:(a.ndim - <NUM_LIT:2>)]):<EOL><INDENT>a_pinv[i] = np.linalg.pinv(<EOL>a[i],<EOL>rcond=rcond if isinstance(rcond, float) else rcond[i])<EOL><DEDENT>return a_pinv<EOL>

Calls `np.linalg.pinv` but corrects its broken batch semantics.


def value_and_gradient(f, xs, use_gradient_tape=False, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [xs]):<EOL><INDENT>is_xs_list_like = isinstance(xs, (tuple, list))<EOL>if not is_xs_list_like:<EOL><INDENT>xs = [xs]<EOL><DEDENT>xs = [tf.convert_to_tensor(value=x, name='<STR_LIT>'.format(i))<EOL>for i, x in enumerate(xs)]<EOL>if tf.executing_eagerly() or use_gradient_tape:<EOL><INDENT>with tf.GradientTape(watch_accessed_variables=False) as tape:<EOL><INDENT>for x in xs:<EOL><INDENT>tape.watch(x)<EOL><DEDENT>y = f(*xs)<EOL><DEDENT>dydx = tape.gradient(y, xs)<EOL><DEDENT>else:<EOL><INDENT>y = f(*xs)<EOL>dydx = tf.gradients(ys=y, xs=xs)<EOL><DEDENT>if not is_xs_list_like:<EOL><INDENT>dydx = dydx[<NUM_LIT:0>]<EOL><DEDENT>return y, dydx<EOL><DEDENT>

Computes `f(*xs)` and its gradients wrt to `*xs`.

    Args:
      f: Python `callable` to be differentiated. If `f` returns a scalar, this
        scalar will be differentiated. If `f` returns a tensor or list of tensors,
        by default a scalar will be computed by adding all their values to produce
        a single scalar. If desired, the tensors can be elementwise multiplied by
        the tensors passed as the `dy` keyword argument to the returned gradient
        function.
      xs: Python list of parameters of f for which to differentiate. (Can also
        be single `Tensor`.)
      use_gradient_tape: Python `bool` indicating that `tf.GradientTape`
        should be used regardless of `tf.executing_eagerly()` status.
        Default value: `False`.
      name: Python `str` name prefixed to ops created by this function.
        Default value: `None` (i.e., `'value_and_gradient'`).

    Returns:
      y: `y = f(*xs)`.
      dydx: Gradient of `y` wrt each of `xs`.


Solves an initial value problem.

        An initial value problem consists of a system of ODEs and an initial
        condition:

        ```none
        dy/dt(t) = ode_fn(t, y(t))
        y(initial_time) = initial_state
        ```

        Here, `t` (also called time) is a scalar `Tensor` of real `dtype` and `y(t)`
        (also called the state at time `t`) is an N-D float or complex `Tensor`.

        Args:
          ode_fn: Function of the form `ode_fn(t, y)`. The input `t` is a scalar
            `Tensor` of real `dtype`. The input `y` and output are both `Tensor`s
            with the same shape and `dtype` as `initial_state`.
          initial_state: N-D float or complex `Tensor` specifying the initial state.
            The `dtype` of `initial_state` must be complex for problems with
            complex-valued states (even if the initial state is real).
          initial_time: Scalar `Tensor` of real `dtype` specifying the initial time.
            Default value: `0.`.
          final_time: Optional scalar `Tensor` of real `dtype` specifying the final
            time to integrate up to. Must satisfy `initial_time < final_time`. If
            unspecified, the solver will use `solution_times[-1]` (see below) as the
            final time. Therefore, at least one of `final_time` or `solution_times`
            must be specified.
            Default value: `None` (i.e., `solution_times[-1]` or error).
          solution_times: Optional 1-D float `Tensor` specifying a list of times. If
            specified, the solver stores the computed state at each of these times
            in the returned `Solution` object. If unspecified, the solver will
            choose this list of times automatically. Must satisfy `initial_time <=
            solution_times[0]`, `solution_times[i] < solution_times[i+1]`, and
            `solution_times[-1] <= final_time` if `final_time` is specified.
            Default value: `None`.
          jacobian_fn: Optional function of the form `jacobian_fn(t, y)`. The input
            `t` is a scalar `Tensor` of real `dtype`. The input `y` has the same
            shape and `dtype` as `initial_state`. The output is a 2N-D `Tensor`
            whose shape is `initial_state.shape + initial_state.shape` and whose
            `dtype` is the same as `initial_state`. In particular, the `(i1, ...,
            iN, j1, ..., jN)`-th entry of `jacobian_fn(t, y)` is the derivative of
            the `(i1, ..., iN)`-th entry of `ode_fn(t, y)` with respect to the `(j1,
            ..., jN)`-th entry of `y`. If this argument is left unspecified, the
            solver will automatically compute the Jacobian if and when it is needed.
            Default value: `None`.
          jacobian_sparsity: Optional 2N-D boolean `Tensor` whose shape is
            `initial_state.shape + initial_state.shape` specifying the sparsity
            pattern of the Jacobian. This argument is ignored if `jacobian_fn` is
            specified.
            Default value: `None`.
          batch_ndims: Optional nonnegative integer. When specified, the first
            `batch_ndims` dimensions of `initial_state` are batch dimensions.
            Default value: `None`.
          previous_results: Optional solver-specific argument used to warm-start
            this invocation of `solve`.
            Default value: `None`.

        Returns:
          solution: Object of type `Solution` containing the solution.
          diagnostics: Object of type `Diagnostics` containing performance
            information.
          results: Solver-specific object which can be used to warm-start the solver
            on a future invocation of `solve`.


pass<EOL>

Number of function evaluations.

        Returns:
          num_ode_fn_evaluations: Scalar `Tensor` of integral `dtype` containing the
            number of function evaluations.


pass<EOL>

Number of Jacobian evaluations.

        Returns:
          num_jacobian_evaluations: Scalar `Tensor` of integral `dtype` containing
            number of Jacobian evaluations.


pass<EOL>

Number of matrix factorizations.

        Returns:
          num_matrix_factorizations: Scalar `Tensor` of integral `dtype` containing
            the number of matrix factorizations.


pass<EOL>

Completion status.

        Returns:
          status: Scalar `Tensor` of integral `dtype` containing the reason for
            termination. -1 on failure, 1 on termination by an event, and 0
            otherwise.


return self.status >= <NUM_LIT:0><EOL>

Boolean indicating whether or not the method succeeded.

        Returns:
          success: Boolean `Tensor` equivalent to `status >= 0`.


with tf.compat.v1.name_scope(name, '<STR_LIT>', [xs, ys]):<EOL><INDENT>if sample_shape is None:<EOL><INDENT>sample_shape = [<NUM_LIT:1>]<EOL><DEDENT>jacobians_diag_res = []<EOL>xs = list(xs) if _is_list_like(xs) else [xs]<EOL>xs = [tf.convert_to_tensor(value=x) for x in xs]<EOL>if not tf.executing_eagerly():<EOL><INDENT>if ys is None:<EOL><INDENT>if fn is None:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>else:<EOL><INDENT>ys = fn(*xs)<EOL><DEDENT><DEDENT>ys = list(ys) if _is_list_like(ys) else [ys]<EOL>if len(xs) != len(ys):<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>for y, x in zip(ys, xs):<EOL><INDENT>y_ = y + tf.zeros_like(x)<EOL>y_ = tf.reshape(y, tf.concat([sample_shape, [-<NUM_LIT:1>]], -<NUM_LIT:1>))<EOL>n = tf.size(input=x) / tf.cast(<EOL>tf.reduce_prod(input_tensor=sample_shape), dtype=tf.int32)<EOL>n = tf.cast(n, dtype=tf.int32)<EOL>loop_vars = [<EOL><NUM_LIT:0>,<EOL>tf.TensorArray(x.dtype, n)<EOL>]<EOL>def loop_body(j):<EOL><INDENT>"""<STR_LIT>"""<EOL>res = tf.gradients(ys=y_[..., j], xs=x)[<NUM_LIT:0>]  <EOL>if res is None:<EOL><INDENT>res = tf.zeros(tf.concat([sample_shape, [<NUM_LIT:1>]], -<NUM_LIT:1>),<EOL>dtype=x.dtype)  <EOL><DEDENT>else:<EOL><INDENT>res = tf.reshape(res, tf.concat([sample_shape, [-<NUM_LIT:1>]], -<NUM_LIT:1>))<EOL>res = tf.expand_dims(res, <NUM_LIT:0>)<EOL>res = res[..., j]<EOL><DEDENT>return res  <EOL><DEDENT>_, jacobian_diag_res = tf.while_loop(<EOL>cond=lambda j, _: j < n,  <EOL>body=lambda j, result: (j + <NUM_LIT:1>, result.write(j, loop_body(j))),<EOL>loop_vars=loop_vars,<EOL>parallel_iterations=parallel_iterations)<EOL>shape_x = tf.shape(input=x)<EOL>reshaped_jacobian_diag = tf.transpose(a=jacobian_diag_res.stack())<EOL>reshaped_jacobian_diag = tf.reshape(reshaped_jacobian_diag, shape_x)<EOL>jacobians_diag_res.append(reshaped_jacobian_diag)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>if fn is None:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if ys is None:<EOL><INDENT>ys = fn(*xs)<EOL><DEDENT>def fn_slice(i, j):<EOL><INDENT>"""<STR_LIT>"""<EOL>def fn_broadcast(*state):<EOL><INDENT>res = fn(*state)<EOL>res = list(res) if _is_list_like(res) else [res]<EOL>if len(res) != len(state):<EOL><INDENT>res *= len(state)<EOL><DEDENT>res = [tf.reshape(r + tf.zeros_like(s),<EOL>tf.concat([sample_shape, [-<NUM_LIT:1>]], -<NUM_LIT:1>))<EOL>for r, s in zip(res, state)]<EOL>return res<EOL><DEDENT>return lambda *state: tf.expand_dims(fn_broadcast(*state)[i], <NUM_LIT:0>)[..., j]<EOL><DEDENT>def make_loop_body(i, x):<EOL><INDENT>"""<STR_LIT>"""<EOL>def _fn(j, result):<EOL><INDENT>res = value_and_gradient(fn_slice(i, j), xs)[<NUM_LIT:1>][i]<EOL>if res is None:<EOL><INDENT>res = tf.zeros(tf.concat([sample_shape, [<NUM_LIT:1>]], -<NUM_LIT:1>), dtype=x.dtype)<EOL><DEDENT>else:<EOL><INDENT>res = tf.reshape(res, tf.concat([sample_shape, [-<NUM_LIT:1>]], -<NUM_LIT:1>))<EOL>res = res[..., j]<EOL><DEDENT>return j + <NUM_LIT:1>, result.write(j, res)<EOL><DEDENT>return _fn<EOL><DEDENT>for i, x in enumerate(xs):<EOL><INDENT>n = tf.size(input=x) / tf.cast(<EOL>tf.reduce_prod(input_tensor=sample_shape), dtype=tf.int32)<EOL>n = tf.cast(n, dtype=tf.int32)<EOL>loop_vars = [<EOL><NUM_LIT:0>,<EOL>tf.TensorArray(x.dtype, n)<EOL>]<EOL>_, jacobian_diag_res = tf.while_loop(<EOL>cond=lambda j, _: j < n,<EOL>body=make_loop_body(i, x),<EOL>loop_vars=loop_vars,<EOL>parallel_iterations=parallel_iterations)<EOL>shape_x = tf.shape(input=x)<EOL>reshaped_jacobian_diag = tf.transpose(a=jacobian_diag_res.stack())<EOL>reshaped_jacobian_diag = tf.reshape(reshaped_jacobian_diag, shape_x)<EOL>jacobians_diag_res.append(reshaped_jacobian_diag)<EOL><DEDENT><DEDENT><DEDENT>return ys, jacobians_diag_res<EOL>

Computes diagonal of the Jacobian matrix of `ys=fn(xs)` wrt `xs`.

      If `ys` is a tensor or a list of tensors of the form `(ys_1, .., ys_n)` and
      `xs` is of the form `(xs_1, .., xs_n)`, the function `jacobians_diag`
      computes the diagonal of the Jacobian matrix, i.e., the partial derivatives
      `(dys_1/dxs_1,.., dys_n/dxs_n`). For definition details, see
      https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant
    #### Example

    ##### Diagonal Hessian of the log-density of a 3D Gaussian distribution

    In this example we sample from a standard univariate normal
    distribution using MALA with `step_size` equal to 0.75.

    ```python
    import tensorflow as tf
    import tensorflow_probability as tfp
    import numpy as np

    tfd = tfp.distributions

    dtype = np.float32
    with tf.Session(graph=tf.Graph()) as sess:
      true_mean = dtype([0, 0, 0])
      true_cov = dtype([[1, 0.25, 0.25], [0.25, 2, 0.25], [0.25, 0.25, 3]])
      chol = tf.linalg.cholesky(true_cov)
      target = tfd.MultivariateNormalTriL(loc=true_mean, scale_tril=chol)

      # Assume that the state is passed as a list of tensors `x` and `y`.
      # Then the target function is defined as follows:
      def target_fn(x, y):
        # Stack the input tensors together
        z = tf.concat([x, y], axis=-1) - true_mean
        return target.log_prob(z)

      sample_shape = [3, 5]
      state = [tf.ones(sample_shape + [2], dtype=dtype),
               tf.ones(sample_shape + [1], dtype=dtype)]
      fn_val, grads = tfp.math.value_and_gradient(target_fn, state)

      # We can either pass the `sample_shape` of the `state` or not, which impacts
      # computational speed of `diag_jacobian`
      _, diag_jacobian_shape_passed = diag_jacobian(
          xs=state, ys=grads, sample_shape=tf.shape(fn_val))
      _, diag_jacobian_shape_none = diag_jacobian(
          xs=state, ys=grads)

      diag_jacobian_shape_passed_ = sess.run(diag_jacobian_shape_passed)
      diag_jacobian_shape_none_ = sess.run(diag_jacobian_shape_none)

    print('hessian computed through `diag_jacobian`, sample_shape passed: ',
          np.concatenate(diag_jacobian_shape_passed_, -1))
    print('hessian computed through `diag_jacobian`, sample_shape skipped',
          np.concatenate(diag_jacobian_shape_none_, -1))

    ```

    Args:
      xs: `Tensor` or a python `list` of `Tensors` of real-like dtypes and shapes
        `sample_shape` + `event_shape_i`, where `event_shape_i` can be different
        for different tensors.
      ys: `Tensor` or a python `list` of `Tensors` of the same dtype as `xs`. Must
          broadcast with the shape of `xs`. Can be omitted if `fn` is provided.
      sample_shape: A common `sample_shape` of the input tensors of `xs`. If not,
        provided, assumed to be `[1]`, which may result in a slow performance of
        `jacobians_diag`.
      fn: Python callable that takes `xs` as an argument (or `*xs`, if it is a
        list) and returns `ys`. Might be skipped if `ys` is provided and
        `tf.enable_eager_execution()` is disabled.
      parallel_iterations: `int` that specifies the allowed number of coordinates
        of the input tensor `xs`, for which the partial derivatives `dys_i/dxs_i`
        can be computed in parallel.
      name: Python `str` name prefixed to `Ops` created by this function.
        Default value: `None` (i.e., "diag_jacobian").

    Returns:
      ys: a list, which coincides with the input `ys`, when provided.
        If the input `ys` is None, `fn(*xs)` gets computed and returned as a list.
      jacobians_diag_res: a `Tensor` or a Python list of `Tensor`s of the same
        dtypes and shapes as the input `xs`. This is the diagonal of the Jacobian
        of ys wrt xs.

    Raises:
      ValueError: if lists `xs` and `ys` have different length or both `ys` and
        `fn` are `None`, or `fn` is None in the eager execution mode.


return isinstance(x, (tuple, list))<EOL>

Helper which returns `True` if input is `list`-like.


with tf.compat.v1.name_scope(name, '<STR_LIT>', [x]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, dtype_hint=tf.float32, name='<STR_LIT:x>')<EOL>dtype = x.dtype.as_numpy_dtype<EOL>eps = np.finfo(dtype).eps.astype(np.float64)<EOL>is_large = tf.abs(x) > (eps**-<NUM_LIT:0.5>).astype(dtype)<EOL>abs_large_x = tf.where(is_large, tf.abs(x), tf.ones_like(x))<EOL>return tf.where(is_large, <NUM_LIT> * tf.math.log(abs_large_x),<EOL>tf.math.log1p(tf.square(x)))<EOL><DEDENT>

Numerically stable calculation of `log(1 + x**2)` for small or large `|x|`.

    For sufficiently large `x` we use the following observation:

    ```none
    log(1 + x**2) =   2 log(|x|) + log(1 + 1 / x**2)
                  --> 2 log(|x|)  as x --> inf
    ```

    Numerically, `log(1 + 1 / x**2)` is `0` when `1 / x**2` is small relative to
    machine epsilon.

    Args:
      x: Float `Tensor` input.
      name: Python string indicating the name of the TensorFlow operation.
        Default value: `'log1psquare'`.

    Returns:
      log1psq: Float `Tensor` representing `log(1. + x**2.)`.


def soft_threshold(x, threshold, name=None):

<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>', [x, threshold]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>threshold = tf.convert_to_tensor(<EOL>value=threshold, dtype=x.dtype, name='<STR_LIT>')<EOL>return tf.sign(x) * tf.maximum(tf.abs(x) - threshold, <NUM_LIT:0.>)<EOL><DEDENT>

Soft Thresholding operator.

    This operator is defined by the equations

    ```none
                                  { x[i] - gamma,  x[i] >   gamma
    SoftThreshold(x, gamma)[i] =  { 0,             x[i] ==  gamma
                                  { x[i] + gamma,  x[i] <  -gamma
    ```

    In the context of proximal gradient methods, we have

    ```none
    SoftThreshold(x, gamma) = prox_{gamma L1}(x)
    ```

    where `prox` is the proximity operator.  Thus the soft thresholding operator
    is used in proximal gradient descent for optimizing a smooth function with
    (non-smooth) L1 regularization, as outlined below.

    The proximity operator is defined as:

    ```none
    prox_r(x) = argmin{ r(z) + 0.5 ||x - z||_2**2 : z },
    ```

    where `r` is a (weakly) convex function, not necessarily differentiable.
    Because the L2 norm is strictly convex, the above argmin is unique.

    One important application of the proximity operator is as follows.  Let `L` be
    a convex and differentiable function with Lipschitz-continuous gradient.  Let
    `R` be a convex lower semicontinuous function which is possibly
    nondifferentiable.  Let `gamma` be an arbitrary positive real.  Then

    ```none
    x_star = argmin{ L(x) + R(x) : x }
    ```

    if and only if the fixed-point equation is satisfied:

    ```none
    x_star = prox_{gamma R}(x_star - gamma grad L(x_star))
    ```

    Proximal gradient descent thus typically consists of choosing an initial value
    `x^{(0)}` and repeatedly applying the update

    ```none
    x^{(k+1)} = prox_{gamma^{(k)} R}(x^{(k)} - gamma^{(k)} grad L(x^{(k)}))
    ```

    where `gamma` is allowed to vary from iteration to iteration.  Specializing to
    the case where `R(x) = ||x||_1`, we minimize `L(x) + ||x||_1` by repeatedly
    applying the update

    ```
    x^{(k+1)} = SoftThreshold(x - gamma grad L(x^{(k)}), gamma)
    ```

    (This idea can also be extended to second-order approximations, although the
    multivariate case does not have a known closed form like above.)

    Args:
      x: `float` `Tensor` representing the input to the SoftThreshold function.
      threshold: nonnegative scalar, `float` `Tensor` representing the radius of
        the interval on which each coordinate of SoftThreshold takes the value
        zero.  Denoted `gamma` above.
      name: Python string indicating the name of the TensorFlow operation.
        Default value: `'soft_threshold'`.

    Returns:
      softthreshold: `float` `Tensor` with the same shape and dtype as `x`,
        representing the value of the SoftThreshold function.

    #### References

    [1]: Yu, Yao-Liang. The Proximity Operator.
         https://www.cs.cmu.edu/~suvrit/teach/yaoliang_proximity.pdf

    [2]: Wikipedia Contributors. Proximal gradient methods for learning.
         _Wikipedia, The Free Encyclopedia_, 2018.
         https://en.wikipedia.org/wiki/Proximal_gradient_methods_for_learning


def clip_by_value_preserve_gradient(t, clip_value_min, clip_value_max,<EOL>name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[t, clip_value_min, clip_value_max]):<EOL><INDENT>t = tf.convert_to_tensor(value=t, name='<STR_LIT:t>')<EOL>clip_t = tf.clip_by_value(t, clip_value_min, clip_value_max)<EOL>return t + tf.stop_gradient(clip_t - t)<EOL><DEDENT>

Clips values to a specified min and max while leaving gradient unaltered.

    Like `tf.clip_by_value`, this function returns a tensor of the same type and
    shape as input `t` but with values clamped to be no smaller than to
    `clip_value_min` and no larger than `clip_value_max`. Unlike
    `tf.clip_by_value`, the gradient is unaffected by this op, i.e.,

    ```python
    tf.gradients(tfp.math.clip_by_value_preserve_gradient(x), x)[0]
    # ==> ones_like(x)
    ```

    Note: `clip_value_min` needs to be smaller or equal to `clip_value_max` for
    correct results.

    Args:
      t: A `Tensor`.
      clip_value_min: A scalar `Tensor`, or a `Tensor` with the same shape
        as `t`. The minimum value to clip by.
      clip_value_max: A scalar `Tensor`, or a `Tensor` with the same shape
        as `t`. The maximum value to clip by.
      name: A name for the operation (optional).
        Default value: `'clip_by_value_preserve_gradient'`.

    Returns:
      clipped_t: A clipped `Tensor`.


def check_results_versus_brute_force(self, x, axis, max_lags, center,<EOL>normalize):

<EOL>axis_len = x.shape[axis]<EOL>if max_lags is None:<EOL><INDENT>max_lags = axis_len - <NUM_LIT:1><EOL><DEDENT>else:<EOL><INDENT>max_lags = min(axis_len - <NUM_LIT:1>, max_lags)<EOL><DEDENT>auto_corr_at_lag = []<EOL>if center:<EOL><INDENT>x -= x.mean(axis=axis, keepdims=True)<EOL><DEDENT>for m in range(max_lags + <NUM_LIT:1>):<EOL><INDENT>auto_corr_at_lag.append(<EOL>(np.take(x, indices=range(<NUM_LIT:0>, axis_len - m), axis=axis) * np.conj(<EOL>np.take(x, indices=range(m, axis_len), axis=axis))).mean(<EOL>axis=axis, keepdims=True))<EOL><DEDENT>rxx = np.concatenate(auto_corr_at_lag, axis=axis)<EOL>if normalize:<EOL><INDENT>rxx /= np.take(rxx, [<NUM_LIT:0>], axis=axis)<EOL><DEDENT>x_ph = tf.compat.v1.placeholder_with_default(<EOL>x, shape=x.shape if self.use_static_shape else None)<EOL>with spectral_ops_test_util.fft_kernel_label_map():<EOL><INDENT>auto_corr = tfp.stats.auto_correlation(<EOL>x_ph,<EOL>axis=axis,<EOL>max_lags=max_lags,<EOL>center=center,<EOL>normalize=normalize)<EOL>if self.use_static_shape:<EOL><INDENT>output_shape = list(x.shape)<EOL>output_shape[axis] = max_lags + <NUM_LIT:1><EOL>self.assertAllEqual(output_shape, auto_corr.shape)<EOL><DEDENT>self.assertAllClose(rxx, self.evaluate(auto_corr), rtol=<NUM_LIT>, atol=<NUM_LIT>)<EOL><DEDENT>

Compute auto-correlation by brute force, then compare to tf result.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[arr, weights, minlength, maxlength,<EOL>axis]):<EOL><INDENT>if axis is None:<EOL><INDENT>return tf.math.bincount(<EOL>arr,<EOL>weights=weights,<EOL>minlength=minlength,<EOL>maxlength=maxlength,<EOL>dtype=dtype)<EOL><DEDENT>arr = tf.convert_to_tensor(value=arr, dtype=tf.int32, name='<STR_LIT>')<EOL>arr_ndims = _get_static_ndims(arr, expect_static=True)<EOL>axis = _make_static_axis_non_negative_list(axis, arr_ndims)<EOL>not_axis = sorted(set(range(arr_ndims)).difference(axis))<EOL>if not not_axis:<EOL><INDENT>return tf.math.bincount(<EOL>arr,<EOL>weights=weights,<EOL>minlength=minlength,<EOL>maxlength=maxlength,<EOL>dtype=dtype)<EOL><DEDENT>flat_arr = _move_dims_to_flat_end(arr, not_axis, arr_ndims, right_end=False)<EOL>if weights is None:<EOL><INDENT>def one_bincount(arr_slice):<EOL><INDENT>return tf.math.bincount(<EOL>arr_slice,<EOL>weights=None,<EOL>minlength=minlength,<EOL>maxlength=maxlength,<EOL>dtype=dtype)<EOL><DEDENT>flat_counts = tf.map_fn(one_bincount, elems=flat_arr, dtype=dtype)<EOL><DEDENT>else:<EOL><INDENT>weights = tf.convert_to_tensor(value=weights, name='<STR_LIT>')<EOL>_get_static_ndims(weights, expect_static=True, expect_ndims=arr_ndims)<EOL>flat_weights = _move_dims_to_flat_end(<EOL>weights, not_axis, arr_ndims, right_end=False)<EOL>def one_bincount(arr_and_weights_slices):<EOL><INDENT>arr_slice, weights_slice = arr_and_weights_slices<EOL>return tf.math.bincount(<EOL>arr_slice,<EOL>weights=weights_slice,<EOL>minlength=minlength,<EOL>maxlength=maxlength,<EOL>dtype=dtype)<EOL><DEDENT>flat_counts = tf.map_fn(<EOL>one_bincount, elems=[flat_arr, flat_weights], dtype=weights.dtype)<EOL><DEDENT>flat_counts_t = tf.transpose(a=flat_counts, perm=[<NUM_LIT:1>, <NUM_LIT:0>])<EOL>_get_static_ndims(flat_counts_t, expect_ndims=<NUM_LIT:2>, expect_static=True)<EOL>not_axis_shape = tf.gather(tf.shape(input=arr), indices=not_axis)<EOL>out_shape = tf.concat([[-<NUM_LIT:1>], not_axis_shape], axis=<NUM_LIT:0>)<EOL>return tf.reshape(flat_counts_t, out_shape)<EOL><DEDENT>

Counts the number of occurrences of each value in an integer array `arr`.

    Works like `tf.math.bincount`, but provides an `axis` kwarg that specifies
    dimensions to reduce over.  With
      `~axis = [i for i in range(arr.ndim) if i not in axis]`,
    this function returns a `Tensor` of shape `[K] + arr.shape[~axis]`.

    If `minlength` and `maxlength` are not given, `K = tf.reduce_max(arr) + 1`
    if `arr` is non-empty, and 0 otherwise.
    If `weights` are non-None, then index `i` of the output stores the sum of the
    value in `weights` at each index where the corresponding value in `arr` is
    `i`.

    Args:
      arr: An `int32` `Tensor` of non-negative values.
      weights: If non-None, must be the same shape as arr. For each value in
        `arr`, the bin will be incremented by the corresponding weight instead of
        1.
      minlength: If given, ensures the output has length at least `minlength`,
        padding with zeros at the end if necessary.
      maxlength: If given, skips values in `arr` that are equal or greater than
        `maxlength`, ensuring that the output has length at most `maxlength`.
      axis: A `0-D` or `1-D` `int32` `Tensor` (with static values) designating
        dimensions in `arr` to reduce over.
        `Default value:` `None`, meaning reduce over all dimensions.
      dtype: If `weights` is None, determines the type of the output bins.
      name: A name scope for the associated operations (optional).

    Returns:
      A vector with the same dtype as `weights` or the given `dtype`. The bin
      values.


<EOL>with tf.compat.v1.name_scope(<EOL>name, default_name='<STR_LIT>', values=[x, edges]):<EOL><INDENT>in_type = dtype_util.common_dtype([x, edges],<EOL>preferred_dtype=tf.float32)<EOL>edges = tf.convert_to_tensor(value=edges, name='<STR_LIT>', dtype=in_type)<EOL>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>', dtype=in_type)<EOL>if (tf.compat.dimension_value(edges.shape[<NUM_LIT:0>]) is not None and<EOL>tf.compat.dimension_value(edges.shape[<NUM_LIT:0>]) < <NUM_LIT:2>):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(edges.shape))<EOL><DEDENT>flattening_x = edges.shape.ndims == <NUM_LIT:1> and x.shape.ndims > <NUM_LIT:1><EOL>if flattening_x:<EOL><INDENT>x_orig_shape = tf.shape(input=x)<EOL>x = tf.reshape(x, [-<NUM_LIT:1>])<EOL><DEDENT>if dtype is None:<EOL><INDENT>dtype = in_type<EOL><DEDENT>dtype = tf.as_dtype(dtype)<EOL>x_permed = distribution_util.rotate_transpose(x, shift=-<NUM_LIT:1>)<EOL>edges_permed = distribution_util.rotate_transpose(edges, shift=-<NUM_LIT:1>)<EOL>searchsorted_type = dtype if dtype in [tf.int32, tf.int64] else None<EOL>almost_output_permed = tf.searchsorted(<EOL>sorted_sequence=edges_permed,<EOL>values=x_permed,<EOL>side='<STR_LIT:right>',<EOL>out_type=searchsorted_type)<EOL>almost_output = tf.cast(<EOL>distribution_util.rotate_transpose(almost_output_permed, shift=<NUM_LIT:1>),<EOL>dtype)<EOL>bins = tf.clip_by_value(almost_output - <NUM_LIT:1>, tf.cast(<NUM_LIT:0>, dtype),<EOL>tf.cast(tf.shape(input=edges)[<NUM_LIT:0>] - <NUM_LIT:2>, dtype))<EOL>if not extend_lower_interval:<EOL><INDENT>low_fill = np.nan if dtype.is_floating else -<NUM_LIT:1><EOL>bins = tf.where(x < tf.expand_dims(edges[<NUM_LIT:0>], <NUM_LIT:0>),<EOL>tf.fill(tf.shape(input=x), tf.cast(low_fill, dtype)),<EOL>bins)<EOL><DEDENT>if not extend_upper_interval:<EOL><INDENT>up_fill = np.nan if dtype.is_floating else tf.shape(input=edges)[<NUM_LIT:0>] - <NUM_LIT:1><EOL>bins = tf.where(x > tf.expand_dims(edges[-<NUM_LIT:1>], <NUM_LIT:0>),<EOL>tf.fill(tf.shape(input=x), tf.cast(up_fill, dtype)), bins)<EOL><DEDENT>if flattening_x:<EOL><INDENT>bins = tf.reshape(bins, x_orig_shape)<EOL><DEDENT>return bins<EOL><DEDENT>

Bin values into discrete intervals.

    Given `edges = [c0, ..., cK]`, defining intervals
    `I0 = [c0, c1)`, `I1 = [c1, c2)`, ..., `I_{K-1} = [c_{K-1}, cK]`,
    This function returns `bins`, such that:
    `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

    Args:
      x:  Numeric `N-D` `Tensor` with `N > 0`.
      edges:  `Tensor` of same `dtype` as `x`.  The first dimension indexes edges
        of intervals.  Must either be `1-D` or have
        `x.shape[1:] == edges.shape[1:]`.  If `rank(edges) > 1`, `edges[k]`
        designates a shape `edges.shape[1:]` `Tensor` of bin edges for the
        corresponding dimensions of `x`.
      extend_lower_interval:  Python `bool`.  If `True`, extend the lowest
        interval `I0` to `(-inf, c1]`.
      extend_upper_interval:  Python `bool`.  If `True`, extend the upper
        interval `I_{K-1}` to `[c_{K-1}, +inf)`.
      dtype: The output type (`int32` or `int64`). `Default value:` `x.dtype`.
        This effects the output values when `x` is below/above the intervals,
        which will be `-1/K+1` for `int` types and `NaN` for `float`s.
        At indices where `x` is `NaN`, the output values will be `0` for `int`
        types and `NaN` for floats.
      name:  A Python string name to prepend to created ops. Default: 'find_bins'

    Returns:
      bins: `Tensor` with same `shape` as `x` and `dtype`.
        Has whole number values.  `bins[i] = k` means the `x[i]` falls into the
        `kth` bin, ie, `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

    Raises:
      ValueError:  If `edges.shape[0]` is determined to be less than 2.

    #### Examples

    Cut a `1-D` array

    ```python
    x = [0., 5., 6., 10., 20.]
    edges = [0., 5., 10.]
    tfp.stats.find_bins(x, edges)
    ==> [0., 0., 1., 1., np.nan]
    ```

    Cut `x` into its deciles

    ```python
    x = tf.random_uniform(shape=(100, 200))
    decile_edges = tfp.stats.quantiles(x, num_quantiles=10)
    bins = tfp.stats.find_bins(x, edges=decile_edges)
    bins.shape
    ==> (100, 200)
    tf.reduce_mean(bins == 0.)
    ==> approximately 0.1
    tf.reduce_mean(bins == 1.)
    ==> approximately 0.1
    ```


with tf.compat.v1.name_scope(name, '<STR_LIT>', values=[x, edges, axis]):<EOL><INDENT>in_dtype = dtype_util.common_dtype([x, edges], preferred_dtype=tf.float32)<EOL>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>', dtype=in_dtype)<EOL>edges = tf.convert_to_tensor(value=edges, name='<STR_LIT>', dtype=in_dtype)<EOL>if axis is None:<EOL><INDENT>x = tf.reshape(x, shape=[-<NUM_LIT:1>])<EOL><DEDENT>else:<EOL><INDENT>x_ndims = _get_static_ndims(<EOL>x, expect_static=True, expect_ndims_at_least=<NUM_LIT:1>)<EOL>axis = _make_static_axis_non_negative_list(axis, x_ndims)<EOL>if not axis:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(axis))<EOL><DEDENT>x = _move_dims_to_flat_end(x, axis, x_ndims, right_end=False)<EOL><DEDENT>bins = find_bins(<EOL>x,<EOL>edges=edges,<EOL>extend_lower_interval=extend_lower_interval,<EOL>extend_upper_interval=extend_upper_interval,<EOL>dtype=tf.int32)<EOL>counts = count_integers(<EOL>bins,<EOL>minlength=tf.shape(input=edges)[<NUM_LIT:0>] - <NUM_LIT:1>,<EOL>maxlength=tf.shape(input=edges)[<NUM_LIT:0>] - <NUM_LIT:1>,<EOL>axis=<NUM_LIT:0>,<EOL>dtype=dtype or in_dtype)<EOL>n_edges = tf.compat.dimension_value(edges.shape[<NUM_LIT:0>])<EOL>if n_edges is not None:<EOL><INDENT>counts.set_shape(<EOL>tf.TensorShape([n_edges - <NUM_LIT:1>]).concatenate(counts.shape[<NUM_LIT:1>:]))<EOL><DEDENT>return counts<EOL><DEDENT>

Count how often `x` falls in intervals defined by `edges`.

    Given `edges = [c0, ..., cK]`, defining intervals
    `I0 = [c0, c1)`, `I1 = [c1, c2)`, ..., `I_{K-1} = [c_{K-1}, cK]`,
    This function counts how often `x` falls into each interval.

    Values of `x` outside of the intervals cause errors.  Consider using
    `extend_lower_interval`, `extend_upper_interval` to deal with this.

    Args:
      x:  Numeric `N-D` `Tensor` with `N > 0`.  If `axis` is not
        `None`, must have statically known number of dimensions. The
        `axis` kwarg determines which dimensions index iid samples.
        Other dimensions of `x` index "events" for which we will compute different
        histograms.
      edges:  `Tensor` of same `dtype` as `x`.  The first dimension indexes edges
        of intervals.  Must either be `1-D` or have `edges.shape[1:]` the same
        as the dimensions of `x` excluding `axis`.
        If `rank(edges) > 1`, `edges[k]` designates a shape `edges.shape[1:]`
        `Tensor` of interval edges for the corresponding dimensions of `x`.
      axis:  Optional `0-D` or `1-D` integer `Tensor` with constant
        values. The axis in `x` that index iid samples.
        `Default value:` `None` (treat every dimension as sample dimension).
      extend_lower_interval:  Python `bool`.  If `True`, extend the lowest
        interval `I0` to `(-inf, c1]`.
      extend_upper_interval:  Python `bool`.  If `True`, extend the upper
        interval `I_{K-1}` to `[c_{K-1}, +inf)`.
      dtype: The output type (`int32` or `int64`). `Default value:` `x.dtype`.
      name:  A Python string name to prepend to created ops.
        `Default value:` 'histogram'

    Returns:
      counts: `Tensor` of type `dtype` and, with
        `~axis = [i for i in range(arr.ndim) if i not in axis]`,
        `counts.shape = [edges.shape[0]] + x.shape[~axis]`.
        With `I` a multi-index into `~axis`, `counts[k][I]` is the number of times
        event(s) fell into the `kth` interval of `edges`.

    #### Examples

    ```python
    # x.shape = [1000, 2]
    # x[:, 0] ~ Uniform(0, 1), x[:, 1] ~ Uniform(1, 2).
    x = tf.stack([tf.random_uniform([1000]), 1 + tf.random_uniform([1000])],
                 axis=-1)

    # edges ==> bins [0, 0.5), [0.5, 1.0), [1.0, 1.5), [1.5, 2.0].
    edges = [0., 0.5, 1.0, 1.5, 2.0]

    tfp.stats.histogram(x, edges)
    ==> approximately [500, 500, 500, 500]

    tfp.stats.histogram(x, edges, axis=0)
    ==> approximately [[500, 500, 0, 0], [0, 0, 500, 500]]
    ```


name = name or '<STR_LIT>'<EOL>allowed_interpolations = {'<STR_LIT>', '<STR_LIT>', '<STR_LIT>', '<STR_LIT>', '<STR_LIT>'}<EOL>if interpolation is None:<EOL><INDENT>interpolation = '<STR_LIT>'<EOL><DEDENT>else:<EOL><INDENT>if interpolation not in allowed_interpolations:<EOL><INDENT>raise ValueError('<STR_LIT>' %<EOL>(allowed_interpolations, interpolation))<EOL><DEDENT><DEDENT>with tf.compat.v1.name_scope(name, values=[x, q]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>if interpolation in {'<STR_LIT>', '<STR_LIT>'} and x.dtype.is_integer:<EOL><INDENT>raise TypeError('<STR_LIT>'.format(<EOL>interpolation, x.dtype))<EOL><DEDENT>q = tf.cast(q, tf.float64)<EOL>_get_static_ndims(q, expect_ndims_no_more_than=<NUM_LIT:1>)<EOL>if validate_args:<EOL><INDENT>q = distribution_util.with_dependencies([<EOL>tf.compat.v1.assert_rank_in(q, [<NUM_LIT:0>, <NUM_LIT:1>]),<EOL>tf.compat.v1.assert_greater_equal(q, tf.cast(<NUM_LIT:0.>, tf.float64)),<EOL>tf.compat.v1.assert_less_equal(q, tf.cast(<NUM_LIT>, tf.float64))<EOL>], q)<EOL><DEDENT>if axis is None:<EOL><INDENT>y = tf.reshape(x, [-<NUM_LIT:1>])<EOL><DEDENT>else:<EOL><INDENT>x_ndims = _get_static_ndims(<EOL>x, expect_static=True, expect_ndims_at_least=<NUM_LIT:1>)<EOL>axis = _make_static_axis_non_negative_list(axis, x_ndims)<EOL>y = _move_dims_to_flat_end(x, axis, x_ndims, right_end=True)<EOL><DEDENT>frac_at_q_or_above = <NUM_LIT:1.> - q / <NUM_LIT><EOL>sorted_y = _sort_tensor(y)<EOL>d = tf.cast(tf.shape(input=y)[-<NUM_LIT:1>], tf.float64)<EOL>def _get_indices(interp_type):<EOL><INDENT>"""<STR_LIT>"""<EOL>if interp_type == '<STR_LIT>':<EOL><INDENT>indices = tf.math.ceil((d - <NUM_LIT:1>) * frac_at_q_or_above)<EOL><DEDENT>elif interp_type == '<STR_LIT>':<EOL><INDENT>indices = tf.floor((d - <NUM_LIT:1>) * frac_at_q_or_above)<EOL><DEDENT>elif interp_type == '<STR_LIT>':<EOL><INDENT>indices = tf.round((d - <NUM_LIT:1>) * frac_at_q_or_above)<EOL><DEDENT>return tf.clip_by_value(<EOL>tf.cast(indices, tf.int32), <NUM_LIT:0>,<EOL>tf.shape(input=y)[-<NUM_LIT:1>] - <NUM_LIT:1>)<EOL><DEDENT>if interpolation in ['<STR_LIT>', '<STR_LIT>', '<STR_LIT>']:<EOL><INDENT>gathered_y = tf.gather(sorted_y, _get_indices(interpolation), axis=-<NUM_LIT:1>)<EOL><DEDENT>elif interpolation == '<STR_LIT>':<EOL><INDENT>gathered_y = <NUM_LIT:0.5> * (<EOL>tf.gather(sorted_y, _get_indices('<STR_LIT>'), axis=-<NUM_LIT:1>) +<EOL>tf.gather(sorted_y, _get_indices('<STR_LIT>'), axis=-<NUM_LIT:1>))<EOL><DEDENT>elif interpolation == '<STR_LIT>':<EOL><INDENT>larger_y_idx = _get_indices('<STR_LIT>')<EOL>exact_idx = (d - <NUM_LIT:1>) * frac_at_q_or_above<EOL>if preserve_gradients:<EOL><INDENT>smaller_y_idx = tf.maximum(larger_y_idx - <NUM_LIT:1>, <NUM_LIT:0>)<EOL>larger_y_idx = tf.minimum(smaller_y_idx + <NUM_LIT:1>, tf.shape(input=y)[-<NUM_LIT:1>] - <NUM_LIT:1>)<EOL>fraction = tf.cast(larger_y_idx, tf.float64) - exact_idx<EOL><DEDENT>else:<EOL><INDENT>smaller_y_idx = _get_indices('<STR_LIT>')<EOL>fraction = tf.math.ceil((d - <NUM_LIT:1>) * frac_at_q_or_above) - exact_idx<EOL><DEDENT>fraction = tf.cast(fraction, y.dtype)<EOL>gathered_y = (<EOL>tf.gather(sorted_y, larger_y_idx, axis=-<NUM_LIT:1>) * (<NUM_LIT:1> - fraction) +<EOL>tf.gather(sorted_y, smaller_y_idx, axis=-<NUM_LIT:1>) * fraction)<EOL><DEDENT>if x.dtype in (tf.bfloat16, tf.float16, tf.float32, tf.float64):<EOL><INDENT>nan_batch_members = tf.reduce_any(<EOL>input_tensor=tf.math.is_nan(x), axis=axis)<EOL>right_rank_matched_shape = tf.pad(<EOL>tensor=tf.shape(input=nan_batch_members),<EOL>paddings=[[<NUM_LIT:0>, tf.rank(input=q)]],<EOL>constant_values=<NUM_LIT:1>)<EOL>nan_batch_members = tf.reshape(<EOL>nan_batch_members, shape=right_rank_matched_shape)<EOL>shape_gathered_y = tf.shape(input=gathered_y)<EOL>nan = np.array(np.nan, gathered_y.dtype.as_numpy_dtype)<EOL>gathered_y = tf.where(<EOL>tf.broadcast_to(nan_batch_members, shape_gathered_y),<EOL>tf.fill(shape_gathered_y, nan),<EOL>gathered_y)<EOL><DEDENT>if keep_dims:<EOL><INDENT>if axis is None:<EOL><INDENT>ones_vec = tf.ones(<EOL>shape=[_get_best_effort_ndims(x) + _get_best_effort_ndims(q)],<EOL>dtype=tf.int32)<EOL>gathered_y *= tf.ones(ones_vec, dtype=x.dtype)<EOL><DEDENT>else:<EOL><INDENT>gathered_y = _insert_back_keep_dims(gathered_y, axis)<EOL><DEDENT><DEDENT>return distribution_util.rotate_transpose(gathered_y, tf.rank(q))<EOL><DEDENT>

Compute the `q`-th percentile(s) of `x`.

    Given a vector `x`, the `q`-th percentile of `x` is the value `q / 100` of the
    way from the minimum to the maximum in a sorted copy of `x`.

    The values and distances of the two nearest neighbors as well as the
    `interpolation` parameter will determine the percentile if the normalized
    ranking does not match the location of `q` exactly.

    This function is the same as the median if `q = 50`, the same as the minimum
    if `q = 0` and the same as the maximum if `q = 100`.

    Multiple percentiles can be computed at once by using `1-D` vector `q`.
    Dimension zero of the returned `Tensor` will index the different percentiles.

    Compare to `numpy.percentile`.

    Args:
      x:  Numeric `N-D` `Tensor` with `N > 0`.  If `axis` is not `None`,
        `x` must have statically known number of dimensions.
      q:  Scalar or vector `Tensor` with values in `[0, 100]`. The percentile(s).
      axis:  Optional `0-D` or `1-D` integer `Tensor` with constant values. The
        axis that index independent samples over which to return the desired
        percentile.  If `None` (the default), treat every dimension as a sample
        dimension, returning a scalar.
      interpolation : {'nearest', 'linear', 'lower', 'higher', 'midpoint'}.
        Default value: 'nearest'.  This specifies the interpolation method to
        use when the desired quantile lies between two data points `i < j`:
          * linear: i + (j - i) * fraction, where fraction is the fractional part
            of the index surrounded by i and j.
          * lower: `i`.
          * higher: `j`.
          * nearest: `i` or `j`, whichever is nearest.
          * midpoint: (i + j) / 2.
        `linear` and `midpoint` interpolation do not work with integer dtypes.
      keep_dims:  Python `bool`. If `True`, the last dimension is kept with size 1
        If `False`, the last dimension is removed from the output shape.
      validate_args:  Whether to add runtime checks of argument validity. If
        False, and arguments are incorrect, correct behavior is not guaranteed.
      preserve_gradients:  Python `bool`.  If `True`, ensure that gradient w.r.t
        the percentile `q` is preserved in the case of linear interpolation.
        If `False`, the gradient will be (incorrectly) zero when `q` corresponds
        to a point in `x`.
      name:  A Python string name to give this `Op`.  Default is 'percentile'

    Returns:
      A `(rank(q) + N - len(axis))` dimensional `Tensor` of same dtype as `x`, or,
        if `axis` is `None`, a `rank(q)` `Tensor`.  The first `rank(q)` dimensions
        index quantiles for different values of `q`.

    Raises:
      ValueError:  If argument 'interpolation' is not an allowed type.
      ValueError:  If interpolation type not compatible with `dtype`.

    #### Examples

    ```python
    # Get 30th percentile with default ('nearest') interpolation.
    x = [1., 2., 3., 4.]
    tfp.stats.percentile(x, q=30.)
    ==> 2.0

    # Get 30th percentile with 'linear' interpolation.
    x = [1., 2., 3., 4.]
    tfp.stats.percentile(x, q=30., interpolation='linear')
    ==> 1.9

    # Get 30th and 70th percentiles with 'lower' interpolation
    x = [1., 2., 3., 4.]
    tfp.stats.percentile(x, q=[30., 70.], interpolation='lower')
    ==> [1., 3.]

    # Get 100th percentile (maximum).  By default, this is computed over every dim
    x = [[1., 2.]
         [3., 4.]]
    tfp.stats.percentile(x, q=100.)
    ==> 4.

    # Treat the leading dim as indexing samples, and find the 100th quantile (max)
    # over all such samples.
    x = [[1., 2.]
         [3., 4.]]
    tfp.stats.percentile(x, q=100., axis=[0])
    ==> [3., 4.]
    ```


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[x, num_quantiles, axis]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>return percentile(<EOL>x,<EOL>q=tf.linspace(<EOL>tf.convert_to_tensor(value=<NUM_LIT:0>, dtype=tf.float64),<EOL>tf.convert_to_tensor(value=<NUM_LIT:100>, dtype=tf.float64),<EOL>num=num_quantiles + <NUM_LIT:1>),<EOL>axis=axis,<EOL>interpolation=interpolation,<EOL>keep_dims=keep_dims,<EOL>validate_args=validate_args,<EOL>preserve_gradients=False)<EOL><DEDENT>

Compute quantiles of `x` along `axis`.

    The quantiles of a distribution are cut points dividing the range into
    intervals with equal probabilities.

    Given a vector `x` of samples, this function estimates the cut points by
    returning `num_quantiles + 1` cut points, `(c0, ..., cn)`, such that, roughly
    speaking, equal number of sample points lie in the `num_quantiles` intervals
    `[c0, c1), [c1, c2), ..., [c_{n-1}, cn]`.  That is,

    * About `1 / n` fraction of the data lies in `[c_{k-1}, c_k)`, `k = 1, ..., n`
    * About `k / n` fraction of the data lies below `c_k`.
    * `c0` is the sample minimum and `cn` is the maximum.

    The exact number of data points in each interval depends on the size of
    `x` (e.g. whether the size is divisible by `n`) and the `interpolation` kwarg.

    Args:
      x:  Numeric `N-D` `Tensor` with `N > 0`.  If `axis` is not `None`,
        `x` must have statically known number of dimensions.
      num_quantiles:  Scalar `integer` `Tensor`.  The number of intervals the
        returned `num_quantiles + 1` cut points divide the range into.
      axis:  Optional `0-D` or `1-D` integer `Tensor` with constant values. The
        axis that index independent samples over which to return the desired
        percentile.  If `None` (the default), treat every dimension as a sample
        dimension, returning a scalar.
      interpolation : {'nearest', 'linear', 'lower', 'higher', 'midpoint'}.
        Default value: 'nearest'.  This specifies the interpolation method to
        use when the fractions `k / n` lie between two data points `i < j`:
          * linear: i + (j - i) * fraction, where fraction is the fractional part
            of the index surrounded by i and j.
          * lower: `i`.
          * higher: `j`.
          * nearest: `i` or `j`, whichever is nearest.
          * midpoint: (i + j) / 2. `linear` and `midpoint` interpolation do not
            work with integer dtypes.
      keep_dims:  Python `bool`. If `True`, the last dimension is kept with size 1
        If `False`, the last dimension is removed from the output shape.
      validate_args:  Whether to add runtime checks of argument validity. If
        False, and arguments are incorrect, correct behavior is not guaranteed.
      name:  A Python string name to give this `Op`.  Default is 'percentile'

    Returns:
      cut_points:  A `rank(x) + 1 - len(axis)` dimensional `Tensor` with same
      `dtype` as `x` and shape `[num_quantiles + 1, ...]` where the trailing shape
      is that of `x` without the dimensions in `axis` (unless `keep_dims is True`)

    Raises:
      ValueError:  If argument 'interpolation' is not an allowed type.
      ValueError:  If interpolation type not compatible with `dtype`.

    #### Examples

    ```python
    # Get quartiles of x with various interpolation choices.
    x = [0.,  1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.]

    tfp.stats.quantiles(x, num_quantiles=4, interpolation='nearest')
    ==> [  0.,   2.,   5.,   8.,  10.]

    tfp.stats.quantiles(x, num_quantiles=4, interpolation='linear')
    ==> [  0. ,   2.5,   5. ,   7.5,  10. ]

    tfp.stats.quantiles(x, num_quantiles=4, interpolation='lower')
    ==> [  0.,   2.,   5.,   7.,  10.]

    # Get deciles of columns of an R x C data set.
    data = load_my_columnar_data(...)
    tfp.stats.quantiles(data, num_quantiles=10)
    ==> Shape [11, C] Tensor
    ```


ndims = x.shape.ndims<EOL>if ndims is None:<EOL><INDENT>shape_const = tf.get_static_value(tf.shape(input=x))<EOL>if shape_const is not None:<EOL><INDENT>ndims = shape_const.ndim<EOL><DEDENT><DEDENT>if ndims is None:<EOL><INDENT>if expect_static:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>' %<EOL>x)<EOL><DEDENT>return<EOL><DEDENT>if expect_ndims is not None:<EOL><INDENT>ndims_message = ('<STR_LIT>'<EOL>% (expect_ndims, x))<EOL>if ndims != expect_ndims:<EOL><INDENT>raise ValueError(ndims_message)<EOL><DEDENT><DEDENT>if expect_ndims_at_least is not None:<EOL><INDENT>ndims_at_least_message = (<EOL>'<STR_LIT>' %<EOL>(expect_ndims_at_least, x))<EOL>if ndims < expect_ndims_at_least:<EOL><INDENT>raise ValueError(ndims_at_least_message)<EOL><DEDENT><DEDENT>if expect_ndims_no_more_than is not None:<EOL><INDENT>ndims_no_more_than_message = (<EOL>'<STR_LIT>' %<EOL>(expect_ndims_no_more_than, x))<EOL>if ndims > expect_ndims_no_more_than:<EOL><INDENT>raise ValueError(ndims_no_more_than_message)<EOL><DEDENT><DEDENT>return ndims<EOL>

Get static number of dimensions and assert that some expectations are met.

    This function returns the number of dimensions 'ndims' of x, as a Python int.

    The optional expect arguments are used to check the ndims of x, but this is
    only done if the static ndims of x is not None.

    Args:
      x:  A Tensor.
      expect_static:  Expect `x` to have statically defined `ndims`.
      expect_ndims:  Optional Python integer.  If provided, assert that x has
        number of dimensions equal to this.
      expect_ndims_no_more_than:  Optional Python integer.  If provided, assert
        that x has no more than this many dimensions.
      expect_ndims_at_least:  Optional Python integer.  If provided, assert that x
        has at least this many dimensions.

    Returns:
      ndims:  A Python integer.

    Raises:
      ValueError:  If any of the expectations above are violated.


ndims_static = _get_static_ndims(<EOL>x,<EOL>expect_ndims=expect_ndims,<EOL>expect_ndims_at_least=expect_ndims_at_least,<EOL>expect_ndims_no_more_than=expect_ndims_no_more_than)<EOL>if ndims_static is not None:<EOL><INDENT>return ndims_static<EOL><DEDENT>return tf.rank(x)<EOL>

Get static ndims if possible.  Fallback on `tf.rank(x)`.


for i in sorted(axis):<EOL><INDENT>x = tf.expand_dims(x, axis=i)<EOL><DEDENT>return x<EOL>

Insert the dims in `axis` back as singletons after being removed.

    Args:
      x:  `Tensor`.
      axis:  Python list of integers.

    Returns:
      `Tensor` with same values as `x`, but additional singleton dimensions.


axis = distribution_util.make_non_negative_axis(axis, ndims)<EOL>axis_const = tf.get_static_value(axis)<EOL>if axis_const is None:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>' %<EOL>axis)<EOL><DEDENT>axis = axis_const + np.zeros([<NUM_LIT:1>], dtype=axis_const.dtype)<EOL>return list(int(dim) for dim in axis)<EOL>

Convert possibly negatively indexed axis to non-negative list of ints.

    Args:
      axis:  Integer Tensor.
      ndims:  Number of dimensions into which axis indexes.

    Returns:
      A list of non-negative Python integers.

    Raises:
      ValueError: If `axis` is not statically defined.


def _move_dims_to_flat_end(x, axis, x_ndims, right_end=True):

if not axis:<EOL><INDENT>return x<EOL><DEDENT>other_dims = sorted(set(range(x_ndims)).difference(axis))<EOL>perm = other_dims + list(axis) if right_end else list(axis) + other_dims<EOL>x_permed = tf.transpose(a=x, perm=perm)<EOL>if x.shape.is_fully_defined():<EOL><INDENT>x_shape = x.shape.as_list()<EOL>other_shape = [x_shape[i] for i in other_dims]<EOL>end_shape = [np.prod([x_shape[i] for i in axis])]<EOL>full_shape = (<EOL>other_shape + end_shape if right_end else end_shape + other_shape)<EOL><DEDENT>else:<EOL><INDENT>other_shape = tf.gather(tf.shape(input=x), other_dims)<EOL>full_shape = tf.concat(<EOL>[other_shape, [-<NUM_LIT:1>]] if right_end else [[-<NUM_LIT:1>], other_shape], axis=<NUM_LIT:0>)<EOL><DEDENT>return tf.reshape(x_permed, shape=full_shape)<EOL>

Move dims corresponding to `axis` in `x` to the end, then flatten.

    Args:
      x: `Tensor` with shape `[B0,B1,...,Bb]`.
      axis:  Python list of indices into dimensions of `x`.
      x_ndims:  Python integer holding number of dimensions in `x`.
      right_end:  Python bool.  Whether to move dims to the right end (else left).

    Returns:
      `Tensor` with value from `x` and dims in `axis` moved to end into one single
        dimension.


sorted_, _ = tf.nn.top_k(tensor, k=tf.shape(input=tensor)[-<NUM_LIT:1>])<EOL>sorted_.set_shape(tensor.shape)<EOL>return sorted_<EOL>

Use `top_k` to sort a `Tensor` along the last dimension.