input_shape = tf.TensorShape(input_shape)<EOL>input_shape = input_shape.with_rank_at_least(<NUM_LIT:2>)<EOL>if tf.compat.dimension_value(input_shape[-<NUM_LIT:1>]) is None:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(input_shape))<EOL><DEDENT>return input_shape[:-<NUM_LIT:1>].concatenate(self.units)<EOL>

Computes the output shape of the layer.

        Args:
          input_shape: Shape tuple (tuple of integers) or list of shape tuples
            (one per output tensor of the layer). Shape tuples can include None for
            free dimensions, instead of an integer.

        Returns:
          output_shape: A tuple representing the output shape.

        Raises:
          ValueError: If innermost dimension of `input_shape` is not defined.


config = {<EOL>'<STR_LIT>': self.units,<EOL>'<STR_LIT>': (tf.keras.activations.serialize(self.activation)<EOL>if self.activation else None),<EOL>'<STR_LIT>':<EOL>tf.keras.initializers.serialize(self.activity_regularizer),<EOL>}<EOL>function_keys = [<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>]<EOL>for function_key in function_keys:<EOL><INDENT>function = getattr(self, function_key)<EOL>if function is None:<EOL><INDENT>function_name = None<EOL>function_type = None<EOL><DEDENT>else:<EOL><INDENT>function_name, function_type = tfp_layers_util.serialize_function(<EOL>function)<EOL><DEDENT>config[function_key] = function_name<EOL>config[function_key + '<STR_LIT>'] = function_type<EOL><DEDENT>base_config = super(_DenseVariational, self).get_config()<EOL>return dict(list(base_config.items()) + list(config.items()))<EOL>

Returns the config of the layer.

        A layer config is a Python dictionary (serializable) containing the
        configuration of a layer. The same layer can be reinstantiated later
        (without its trained weights) from this configuration.

        Returns:
          config: A Python dictionary of class keyword arguments and their
            serialized values.


config = config.copy()<EOL>function_keys = [<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>'<STR_LIT>',<EOL>]<EOL>for function_key in function_keys:<EOL><INDENT>serial = config[function_key]<EOL>function_type = config.pop(function_key + '<STR_LIT>')<EOL>if serial is not None:<EOL><INDENT>config[function_key] = tfp_layers_util.deserialize_function(<EOL>serial,<EOL>function_type=function_type)<EOL><DEDENT><DEDENT>return cls(**config)<EOL>

Creates a layer from its config.

        This method is the reverse of `get_config`, capable of instantiating the
        same layer from the config dictionary.

        Args:
          config: A Python dictionary, typically the output of `get_config`.

        Returns:
          layer: A layer instance.


Construct layer.

        Args:
          ${args}


Construct layer.

        Args:
          ${args}
          seed: Python scalar `int` which initializes the random number
            generator. Default value: `None` (i.e., use global seed).


config = {<EOL>'<STR_LIT>': self.seed,<EOL>}<EOL>base_config = super(DenseFlipout, self).get_config()<EOL>return dict(list(base_config.items()) + list(config.items()))<EOL>


if '<STR_LIT>' not in kwargs and '<STR_LIT>' in kwargs:<EOL><INDENT>kwargs['<STR_LIT>'] = (kwargs.pop('<STR_LIT>'),)<EOL><DEDENT>super(DenseVariational, self).__init__(<EOL>activity_regularizer=tf.keras.regularizers.get(activity_regularizer),<EOL>**kwargs)<EOL>self.units = int(units)<EOL>self._make_posterior_fn = make_posterior_fn<EOL>self._make_prior_fn = make_prior_fn<EOL>self._kl_divergence_fn = _make_kl_divergence_penalty(<EOL>kl_use_exact, weight=kl_weight)<EOL>self.activation = tf.keras.activations.get(activation)<EOL>self.use_bias = use_bias<EOL>self.supports_masking = False<EOL>self.input_spec = tf.keras.layers.InputSpec(min_ndim=<NUM_LIT:2>)<EOL>

Creates the `DenseVariational` layer.

        Arguments:
          units: Positive integer, dimensionality of the output space.
          make_posterior_fn: Python callable taking `tf.size(kernel)`,
            `tf.size(bias)`, `dtype` and returns another callable which takes an
            input and produces a `tfd.Distribution` instance.
          make_prior_fn: Python callable taking `tf.size(kernel)`, `tf.size(bias)`,
            `dtype` and returns another callable which takes an input and produces a
            `tfd.Distribution` instance.
          kl_weight: Amount by which to scale the KL divergence loss between prior
            and posterior.
          kl_use_exact: Python `bool` indicating that the analytical KL divergence
            should be used rather than a Monte Carlo approximation.
          activation: Activation function to use.
            If you don't specify anything, no activation is applied
            (ie. "linear" activation: `a(x) = x`).
          use_bias: Boolean, whether the layer uses a bias vector.
          activity_regularizer: Regularizer function applied to
            the output of the layer (its "activation")..
          **kwargs: Extra arguments forwarded to `tf.keras.layers.Layer`.


super(AutoregressiveTransform, self).__init__(self._transform, **kwargs)<EOL>if made.params != <NUM_LIT:2>:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'.format(made.params))<EOL><DEDENT>self._made = made<EOL>

Constructs the AutoregressiveTransform layer.

        Arguments:
          made: A `Made` layer, which must output two parameters for each input.
          **kwargs: Additional keyword arguments passed to `tf.keras.Layer`.


super(VariableLayer, self).__init__(**kwargs)<EOL>self.activation = tf.keras.activations.get(activation)<EOL>self.initializer = tf.keras.initializers.get(initializer)<EOL>self.regularizer = tf.keras.regularizers.get(regularizer)<EOL>self.constraint = tf.keras.constraints.get(constraint)<EOL>shape = tf.get_static_value(shape)<EOL>if shape is None:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>shape = np.array(shape, dtype=np.int32)<EOL>ndims = len(shape.shape)<EOL>if ndims > <NUM_LIT:1>:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>shape = shape.reshape(-<NUM_LIT:1>)  <EOL>self._var = self.add_weight(<EOL>'<STR_LIT>',<EOL>shape=shape,<EOL>initializer=self.initializer,<EOL>regularizer=self.regularizer,<EOL>constraint=self.constraint,<EOL>dtype=dtype,<EOL>trainable=kwargs.get('<STR_LIT>', True))<EOL>

Creates the `VariableLayer`.

        Arguments:
          shape: integer or integer vector specifying the shape of the output of
            this layer.
          dtype: TensorFlow `dtype` of the variable created by this layer.
            Default value: `None` (i.e., `tf.as_dtype(tf.keras.backend.floatx())`).
          activation: Activation function to use.  If you don't specify anything, no
            activation is applied (ie. "linear" activation: `a(x) = x`).
            Default value: `None`.
          initializer: Initializer for the `constant` vector. For example, to
            initialize a trainable (initially standard) `Normal` with
            `tf.math.softplus` transformed scale, one might use:
            ```python
            tfp.layers.BlockwiseInitializer([
                'zeros',
                tf.keras.initializers.Constant(np.log(np.expm1(1.))),  # = 0.541325
            ], sizes=[1, 1])
            ```
            Default value: `'zeros'`.
          regularizer: Regularizer function applied to the `constant` vector.
            Default value: `None`.
          constraint: Constraint function applied to the `constant` vector.
            Default value: `None`.
          **kwargs: Extra arguments forwarded to `tf.keras.layers.Layer`.


def _grad_neg_log_likelihood_and_fim(model_matrix, linear_response, response,<EOL>model):

<EOL>mean, variance, grad_mean = model(linear_response)<EOL>is_valid = (<EOL>tf.math.is_finite(grad_mean) & tf.not_equal(grad_mean, <NUM_LIT:0.>)<EOL>& tf.math.is_finite(variance) & (variance > <NUM_LIT:0.>))<EOL>def _mask_if_invalid(x, mask):<EOL><INDENT>mask = tf.fill(<EOL>tf.shape(input=x), value=np.array(mask, x.dtype.as_numpy_dtype))<EOL>return tf.where(is_valid, x, mask)<EOL><DEDENT>v = (response - mean) * _mask_if_invalid(grad_mean, <NUM_LIT:1>) / _mask_if_invalid(<EOL>variance, np.inf)<EOL>grad_log_likelihood = sparse_or_dense_matvecmul(<EOL>model_matrix, v, adjoint_a=True)<EOL>fim_middle = _mask_if_invalid(grad_mean, <NUM_LIT:0.>)**<NUM_LIT:2> / _mask_if_invalid(<EOL>variance, np.inf)<EOL>return -grad_log_likelihood, fim_middle<EOL>

Computes the neg-log-likelihood gradient and Fisher information for a GLM.

    Note that Fisher information is related to the Hessian of the log-likelihood
    by the equation

    ```none
    FisherInfo = E[Hessian with respect to model_coefficients of -LogLikelihood(
        Y | model_matrix, model_coefficients)]
    ```

    where `LogLikelihood` is the log-likelihood of a generalized linear model
    parameterized by `model_matrix` and `model_coefficients`, and the expectation
    is taken over Y, distributed according to the same GLM with the same parameter
    values.

    Args:
      model_matrix: (Batch of) matrix-shaped, `float` `Tensor` or `SparseTensor`
        where each row represents a sample's features.  Has shape `[N, n]` where
        `N` is the number of data samples and `n` is the number of features per
        sample.
      linear_response: (Batch of) vector-shaped `Tensor` with the same dtype as
        `model_matrix`, equal to `model_matix @ model_coefficients` where
        `model_coefficients` are the coefficients of the linear component of the
        GLM.
      response: (Batch of) vector-shaped `Tensor` with the same dtype as
        `model_matrix` where each element represents a sample's observed response
        (to the corresponding row of features).
      model: `tfp.glm.ExponentialFamily`-like instance, which specifies the link
        function and distribution of the GLM, and thus characterizes the negative
        log-likelihood. Must have sufficient statistic equal to the response, that
        is, `T(y) = y`.

    Returns:
      grad_neg_log_likelihood: (Batch of) vector-shaped `Tensor` with the same
        shape and dtype as a single row of `model_matrix`, representing the
        gradient of the negative log likelihood of `response` given linear
        response `linear_response`.
      fim_middle: (Batch of) vector-shaped `Tensor` with the same shape and dtype
        as a single column of `model_matrix`, satisfying the equation
        `Fisher information =
        Transpose(model_matrix)
        @ diag(fim_middle)
        @ model_matrix`.


One step of (the outer loop of) the GLM fitting algorithm.

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

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

    ```none
    model_coefficients_update^* = argmin{
        -LogLikelihood(model_coefficients_start + model_coefficients_update')
          + l1_regularizer *
              ||model_coefficients_start + model_coefficients_update'||_1
          + l2_regularizer *
              ||model_coefficients_start + model_coefficients_update'||_2**2
        : model_coefficients_update' }
    ```

    where in each iteration `model_coefficients_update'` has at most one nonzero
    coordinate.

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

    Args:
      model_matrix: (Batch of) matrix-shaped, `float` `Tensor` or `SparseTensor`
        where each row represents a sample's features.  Has shape `[N, n]` where
        `N` is the number of data samples and `n` is the number of features per
        sample.
      response: (Batch of) vector-shaped `Tensor` with the same dtype as
        `model_matrix` where each element represents a sample's observed response
        (to the corresponding row of features).
      model: `tfp.glm.ExponentialFamily`-like instance, which specifies the link
        function and distribution of the GLM, and thus characterizes the negative
        log-likelihood which will be minimized. Must have sufficient statistic
        equal to the response, that is, `T(y) = y`.
      model_coefficients_start: (Batch of) vector-shaped, `float` `Tensor` with
        the same dtype as `model_matrix`, representing the initial values of the
        coefficients for the GLM regression.  Has shape `[n]` where `model_matrix`
        has shape `[N, n]`.
      tolerance: scalar, `float` `Tensor` representing the convergence threshold.
        The optimization step will terminate early, returning its current value of
        `model_coefficients_start + model_coefficients_update`, once the following
        condition is met:
        `||model_coefficients_update_end - model_coefficients_update_start||_2
           / (1 + ||model_coefficients_start||_2)
         < sqrt(tolerance)`,
        where `model_coefficients_update_end` is the value of
        `model_coefficients_update` at the end of a sweep and
        `model_coefficients_update_start` is the value of
        `model_coefficients_update` at the beginning of that sweep.
      l1_regularizer: scalar, `float` `Tensor` representing the weight of the L1
        regularization term (see equation above).
      l2_regularizer: scalar, `float` `Tensor` representing the weight of the L2
        regularization term (see equation above).
        Default value: `None` (i.e., no L2 regularization).
      maximum_full_sweeps: Python integer specifying maximum number of sweeps to
        run.  A "sweep" consists of an iteration of coordinate descent on each
        coordinate. After this many sweeps, the algorithm will terminate even if
        convergence has not been reached.
        Default value: `1`.
      learning_rate: scalar, `float` `Tensor` representing a multiplicative factor
        used to dampen the proximal gradient descent steps.
        Default value: `None` (i.e., factor is conceptually `1`).
      name: Python string representing the name of the TensorFlow operation. The
        default name is `"fit_sparse_one_step"`.

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


graph_deps = [<EOL>model_matrix,<EOL>response,<EOL>model_coefficients_start,<EOL>l1_regularizer,<EOL>l2_regularizer,<EOL>maximum_iterations,<EOL>maximum_full_sweeps_per_iteration,<EOL>tolerance,<EOL>learning_rate,<EOL>]<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>', graph_deps):<EOL><INDENT>def _grad_neg_log_likelihood_and_fim_fn(x):<EOL><INDENT>predicted_linear_response = sparse_or_dense_matvecmul(model_matrix, x)<EOL>g, h_middle = _grad_neg_log_likelihood_and_fim(<EOL>model_matrix, predicted_linear_response, response, model)<EOL>return g, model_matrix, h_middle<EOL><DEDENT>return tfp.optimizer.proximal_hessian_sparse_minimize(<EOL>_grad_neg_log_likelihood_and_fim_fn,<EOL>x_start=model_coefficients_start,<EOL>l1_regularizer=l1_regularizer,<EOL>l2_regularizer=l2_regularizer,<EOL>maximum_iterations=maximum_iterations,<EOL>maximum_full_sweeps_per_iteration=maximum_full_sweeps_per_iteration,<EOL>learning_rate=learning_rate,<EOL>tolerance=tolerance,<EOL>name=name)<EOL><DEDENT>

r"""Fits a GLM using coordinate-wise FIM-informed proximal gradient descent.

    This function uses a L1- and L2-regularized, second-order quasi-Newton method
    to find maximum-likelihood parameters for the given model and observed data.
    The second-order approximations use negative Fisher information in place of
    the Hessian, that is,

    ```none
    FisherInfo = E_Y[Hessian with respect to model_coefficients of -LogLikelihood(
        Y | model_matrix, current value of model_coefficients)]
    ```

    For large, sparse data sets, `model_matrix` should be supplied as a
    `SparseTensor`.

    Args:
      model_matrix: (Batch of) matrix-shaped, `float` `Tensor` or `SparseTensor`
        where each row represents a sample's features.  Has shape `[N, n]` where
        `N` is the number of data samples and `n` is the number of features per
        sample.
      response: (Batch of) vector-shaped `Tensor` with the same dtype as
        `model_matrix` where each element represents a sample's observed response
        (to the corresponding row of features).
      model: `tfp.glm.ExponentialFamily`-like instance, which specifies the link
        function and distribution of the GLM, and thus characterizes the negative
        log-likelihood which will be minimized. Must have sufficient statistic
        equal to the response, that is, `T(y) = y`.
      model_coefficients_start: (Batch of) vector-shaped, `float` `Tensor` with
        the same dtype as `model_matrix`, representing the initial values of the
        coefficients for the GLM regression.  Has shape `[n]` where `model_matrix`
        has shape `[N, n]`.
      tolerance: scalar, `float` `Tensor` representing the tolerance for each
        optiization step; see the `tolerance` argument of `fit_sparse_one_step`.
      l1_regularizer: scalar, `float` `Tensor` representing the weight of the L1
        regularization term.
      l2_regularizer: scalar, `float` `Tensor` representing the weight of the L2
        regularization term.
        Default value: `None` (i.e., no L2 regularization).
      maximum_iterations: Python integer specifying maximum number of iterations
        of the outer loop of the optimizer (i.e., maximum number of calls to
        `fit_sparse_one_step`).  After this many iterations of the outer loop, the
        algorithm will terminate even if the return value `model_coefficients` has
        not converged.
        Default value: `1`.
      maximum_full_sweeps_per_iteration: Python integer specifying the maximum
        number of coordinate descent sweeps allowed in each iteration.
        Default value: `1`.
      learning_rate: scalar, `float` `Tensor` representing a multiplicative factor
        used to dampen the proximal gradient descent steps.
        Default value: `None` (i.e., factor is conceptually `1`).
      name: Python string representing the name of the TensorFlow operation.
        The default name is `"fit_sparse"`.

    Returns:
      model_coefficients: (Batch of) `Tensor` of the same shape and dtype as
        `model_coefficients_start`, representing the computed model coefficients
        which minimize the regularized negative log-likelihood.
      is_converged: scalar, `bool` `Tensor` indicating whether the minimization
        procedure converged across all batches within the specified number of
        iterations.  Here convergence means that an iteration of the inner loop
        (`fit_sparse_one_step`) returns `True` for its `is_converged` output
        value.
      iter: scalar, `int` `Tensor` indicating the actual number of iterations of
        the outer loop of the optimizer completed (i.e., number of calls to
        `fit_sparse_one_step` before achieving convergence).

    #### Example

    ```python
    from __future__ import print_function
    import numpy as np
    import tensorflow as tf
    import tensorflow_probability as tfp
    tfd = tfp.distributions

    def make_dataset(n, d, link, scale=1., dtype=np.float32):
      model_coefficients = tfd.Uniform(
          low=np.array(-1, dtype), high=np.array(1, dtype)).sample(
              d, seed=42)
      radius = np.sqrt(2.)
      model_coefficients *= radius / tf.linalg.norm(model_coefficients)
      mask = tf.random_shuffle(tf.range(d)) < tf.to_int32(0.5 * tf.to_float(d))
      model_coefficients = tf.where(mask, model_coefficients,
                                    tf.zeros_like(model_coefficients))
      model_matrix = tfd.Normal(
          loc=np.array(0, dtype), scale=np.array(1, dtype)).sample(
              [n, d], seed=43)
      scale = tf.convert_to_tensor(scale, dtype)
      linear_response = tf.matmul(model_matrix,
                                  model_coefficients[..., tf.newaxis])[..., 0]
      if link == 'linear':
        response = tfd.Normal(loc=linear_response, scale=scale).sample(seed=44)
      elif link == 'probit':
        response = tf.cast(
            tfd.Normal(loc=linear_response, scale=scale).sample(seed=44) > 0,
                       dtype)
      elif link == 'logit':
        response = tfd.Bernoulli(logits=linear_response).sample(seed=44)
      else:
        raise ValueError('unrecognized true link: {}'.format(link))
      return model_matrix, response, model_coefficients, mask

    with tf.Session() as sess:
      x_, y_, model_coefficients_true_, _ = sess.run(make_dataset(
          n=int(1e5), d=100, link='probit'))

      model = tfp.glm.Bernoulli()
      model_coefficients_start = tf.zeros(x_.shape[-1], np.float32)

      model_coefficients, is_converged, num_iter = tfp.glm.fit_sparse(
          model_matrix=tf.convert_to_tensor(x_),
          response=tf.convert_to_tensor(y_),
          model=model,
          model_coefficients_start=model_coefficients_start,
          l1_regularizer=800.,
          l2_regularizer=None,
          maximum_iterations=10,
          maximum_full_sweeps_per_iteration=10,
          tolerance=1e-6,
          learning_rate=None)

      model_coefficients_, is_converged_, num_iter_ = sess.run([
          model_coefficients, is_converged, num_iter])

      print("is_converged:", is_converged_)
      print("    num_iter:", num_iter_)
      print("\nLearned / True")
      print(np.concatenate(
          [[model_coefficients_], [model_coefficients_true_]], axis=0).T)

    # ==>
    # is_converged: True
    #     num_iter: 1
    #
    # Learned / True
    # [[ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [ 0.11195257  0.12484948]
    #  [ 0.          0.        ]
    #  [ 0.05191106  0.06394956]
    #  [-0.15090358 -0.15325639]
    #  [-0.18187316 -0.18825999]
    #  [-0.06140942 -0.07994166]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [ 0.14474444  0.15810856]
    #  [ 0.          0.        ]
    #  [-0.25249591 -0.24260855]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [-0.03888761 -0.06755984]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [-0.0192222  -0.04169233]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [ 0.01434913  0.03568212]
    #  [-0.11336883 -0.12873614]
    #  [ 0.          0.        ]
    #  [-0.24496339 -0.24048163]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [ 0.04088281  0.06565224]
    #  [-0.12784363 -0.13359821]
    #  [ 0.05618424  0.07396613]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [ 0.         -0.01719233]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [-0.00076072 -0.03607186]
    #  [ 0.21801499  0.21146794]
    #  [-0.02161094 -0.04031265]
    #  [ 0.0918689   0.10487888]
    #  [ 0.0106154   0.03233612]
    #  [-0.07817317 -0.09725142]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [-0.23725343 -0.24194022]
    #  [ 0.          0.        ]
    #  [-0.08725718 -0.1048776 ]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [-0.02114314 -0.04145789]
    #  [ 0.          0.        ]
    #  [ 0.          0.        ]
    #  [-0.02710908 -0.04590397]
    #  [ 0.15293184  0.15415154]
    #  [ 0.2114463   0.2088728 ]
    #  [-0.10969634 -0.12368613]
    #  [ 0.         -0.01505797]
    #  [-0.01140458 -0.03234904]
    #  [ 0.16051085  0.1680062 ]
    #  [ 0.09816848  0.11094204]
    ```

    #### References

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

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


mean = tf.convert_to_tensor(value=mean, name='<STR_LIT>')<EOL>s2 = np.log(<NUM_LIT>).astype(mean.dtype.as_numpy_dtype())<EOL>return tfd.LogNormal(tf.math.log(mean) - <NUM_LIT:0.5> * s2, np.sqrt(s2))<EOL>

Helper which creates a LogNormal with a specific mean.


graph_deps = [model_matrix, response, model_coefficients_start,<EOL>predicted_linear_response_start, dispersion, offset,<EOL>learning_rate, maximum_iterations]<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>', graph_deps):<EOL><INDENT>[<EOL>model_matrix,<EOL>response,<EOL>model_coefficients_start,<EOL>predicted_linear_response_start,<EOL>offset,<EOL>] = prepare_args(<EOL>model_matrix,<EOL>response,<EOL>model_coefficients_start,<EOL>predicted_linear_response_start,<EOL>offset)<EOL>if convergence_criteria_fn is None:<EOL><INDENT>convergence_criteria_fn = (<EOL>convergence_criteria_small_relative_norm_weights_change())<EOL><DEDENT>def _body(<EOL>is_converged_previous,<EOL>iter_,<EOL>model_coefficients_previous,<EOL>predicted_linear_response_previous):<EOL><INDENT>"""<STR_LIT>"""<EOL>model_coefficients_next, predicted_linear_response_next = fit_one_step(<EOL>model_matrix,<EOL>response,<EOL>model,<EOL>model_coefficients_previous,<EOL>predicted_linear_response_previous,<EOL>l2_regularizer,<EOL>dispersion,<EOL>offset,<EOL>learning_rate,<EOL>fast_unsafe_numerics)<EOL>is_converged_next = convergence_criteria_fn(<EOL>is_converged_previous=is_converged_previous,<EOL>iter_=iter_,<EOL>model_coefficients_previous=model_coefficients_previous,<EOL>predicted_linear_response_previous=predicted_linear_response_previous,<EOL>model_coefficients_next=model_coefficients_next,<EOL>predicted_linear_response_next=predicted_linear_response_next,<EOL>response=response,<EOL>model=model,<EOL>dispersion=dispersion)<EOL>return [<EOL>is_converged_next,<EOL>iter_ + <NUM_LIT:1>,<EOL>model_coefficients_next,<EOL>predicted_linear_response_next,<EOL>]<EOL><DEDENT>[<EOL>is_converged,<EOL>iter_,<EOL>model_coefficients,<EOL>predicted_linear_response,<EOL>] = tf.while_loop(<EOL>cond=lambda is_converged, *args: tf.logical_not(is_converged),<EOL>body=_body,<EOL>loop_vars=[<EOL>tf.zeros([], np.bool),   <EOL>tf.zeros([], np.int32),  <EOL>model_coefficients_start,<EOL>predicted_linear_response_start,<EOL>],<EOL>maximum_iterations=maximum_iterations)<EOL>return [<EOL>model_coefficients,<EOL>predicted_linear_response,<EOL>is_converged,<EOL>iter_<EOL>]<EOL><DEDENT>

Runs multiple Fisher scoring steps.

    Args:
      model_matrix: (Batch of) `float`-like, matrix-shaped `Tensor` where each row
        represents a sample's features.
      response: (Batch of) vector-shaped `Tensor` where each element represents a
        sample's observed response (to the corresponding row of features). Must
        have same `dtype` as `model_matrix`.
      model: `tfp.glm.ExponentialFamily`-like instance which implicitly
        characterizes a negative log-likelihood loss by specifying the
        distribuion's `mean`, `gradient_mean`, and `variance`.
      model_coefficients_start: Optional (batch of) vector-shaped `Tensor`
        representing the initial model coefficients, one for each column in
        `model_matrix`. Must have same `dtype` as `model_matrix`.
        Default value: Zeros.
      predicted_linear_response_start: Optional `Tensor` with `shape`, `dtype`
        matching `response`; represents `offset` shifted initial linear
        predictions based on `model_coefficients_start`.
        Default value: `offset` if `model_coefficients is None`, and
        `tf.linalg.matvec(model_matrix, model_coefficients_start) + offset`
        otherwise.
      l2_regularizer: Optional scalar `Tensor` representing L2 regularization
        penalty, i.e.,
        `loss(w) = sum{-log p(y[i]|x[i],w) : i=1..n} + l2_regularizer ||w||_2^2`.
        Default value: `None` (i.e., no L2 regularization).
      dispersion: Optional (batch of) `Tensor` representing `response` dispersion,
        i.e., as in, `p(y|theta) := exp((y theta - A(theta)) / dispersion)`.
        Must broadcast with rows of `model_matrix`.
        Default value: `None` (i.e., "no dispersion").
      offset: Optional `Tensor` representing constant shift applied to
        `predicted_linear_response`.  Must broadcast to `response`.
        Default value: `None` (i.e., `tf.zeros_like(response)`).
      convergence_criteria_fn: Python `callable` taking:
        `is_converged_previous`, `iter_`, `model_coefficients_previous`,
        `predicted_linear_response_previous`, `model_coefficients_next`,
        `predicted_linear_response_next`, `response`, `model`, `dispersion` and
        returning a `bool` `Tensor` indicating that Fisher scoring has converged.
        See `convergence_criteria_small_relative_norm_weights_change` as an
        example function.
        Default value: `None` (i.e.,
        `convergence_criteria_small_relative_norm_weights_change`).
      learning_rate: Optional (batch of) scalar `Tensor` used to dampen iterative
        progress. Typically only needed if optimization diverges, should be no
        larger than `1` and typically very close to `1`.
        Default value: `None` (i.e., `1`).
      fast_unsafe_numerics: Optional Python `bool` indicating if faster, less
        numerically accurate methods can be employed for computing the weighted
        least-squares solution.
        Default value: `True` (i.e., "fast but possibly diminished accuracy").
      maximum_iterations: Optional maximum number of iterations of Fisher scoring
        to run; "and-ed" with result of `convergence_criteria_fn`.
        Default value: `None` (i.e., `infinity`).
      name: Python `str` used as name prefix to ops created by this function.
        Default value: `"fit"`.

    Returns:
      model_coefficients: (Batch of) vector-shaped `Tensor`; represents the
        fitted model coefficients, one for each column in `model_matrix`.
      predicted_linear_response: `response`-shaped `Tensor` representing linear
        predictions based on new `model_coefficients`, i.e.,
        `tf.linalg.matvec(model_matrix, model_coefficients) + offset`.
      is_converged: `bool` `Tensor` indicating that the returned
        `model_coefficients` met the `convergence_criteria_fn` criteria within the
        `maximum_iterations` limit.
      iter_: `int32` `Tensor` indicating the number of iterations taken.

    #### Example

    ```python
    from __future__ import print_function
    import numpy as np
    import tensorflow as tf
    import tensorflow_probability as tfp
    tfd = tfp.distributions

    def make_dataset(n, d, link, scale=1., dtype=np.float32):
      model_coefficients = tfd.Uniform(
          low=np.array(-1, dtype),
          high=np.array(1, dtype)).sample(d, seed=42)
      radius = np.sqrt(2.)
      model_coefficients *= radius / tf.linalg.norm(model_coefficients)
      model_matrix = tfd.Normal(
          loc=np.array(0, dtype),
          scale=np.array(1, dtype)).sample([n, d], seed=43)
      scale = tf.convert_to_tensor(scale, dtype)
      linear_response = tf.tensordot(
          model_matrix, model_coefficients, axes=[[1], [0]])
      if link == 'linear':
        response = tfd.Normal(loc=linear_response, scale=scale).sample(seed=44)
      elif link == 'probit':
        response = tf.cast(
            tfd.Normal(loc=linear_response, scale=scale).sample(seed=44) > 0,
            dtype)
      elif link == 'logit':
        response = tfd.Bernoulli(logits=linear_response).sample(seed=44)
      else:
        raise ValueError('unrecognized true link: {}'.format(link))
      return model_matrix, response, model_coefficients

    X, Y, w_true = make_dataset(n=int(1e6), d=100, link='probit')

    w, linear_response, is_converged, num_iter = tfp.glm.fit(
        model_matrix=X,
        response=Y,
        model=tfp.glm.BernoulliNormalCDF())
    log_likelihood = tfp.glm.BernoulliNormalCDF().log_prob(Y, linear_response)

    with tf.Session() as sess:
      [w_, linear_response_, is_converged_, num_iter_, Y_, w_true_,
       log_likelihood_] = sess.run([
          w, linear_response, is_converged, num_iter, Y, w_true,
          log_likelihood])

    print('is_converged: ', is_converged_)
    print('    num_iter: ', num_iter_)
    print('    accuracy: ', np.mean((linear_response_ > 0.) == Y_))
    print('    deviance: ', 2. * np.mean(log_likelihood_))
    print('||w0-w1||_2 / (1+||w0||_2): ', (np.linalg.norm(w_true_ - w_, ord=2) /
                                           (1. + np.linalg.norm(w_true_, ord=2))))

    # ==>
    # is_converged:  True
    #     num_iter:  6
    #     accuracy:  0.804382
    #     deviance:  -0.820746600628
    # ||w0-w1||_2 / (1+||w0||_2):  0.00619245105309
    ```


graph_deps = [model_matrix, response, model_coefficients_start,<EOL>predicted_linear_response_start, dispersion, learning_rate]<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>', graph_deps):<EOL><INDENT>[<EOL>model_matrix,<EOL>response,<EOL>model_coefficients_start,<EOL>predicted_linear_response_start,<EOL>offset,<EOL>] = prepare_args(<EOL>model_matrix,<EOL>response,<EOL>model_coefficients_start,<EOL>predicted_linear_response_start,<EOL>offset)<EOL>mean, variance, grad_mean = model(predicted_linear_response_start)<EOL>is_valid = (<EOL>tf.math.is_finite(grad_mean) & tf.not_equal(grad_mean, <NUM_LIT:0.>)<EOL>& tf.math.is_finite(variance) & (variance > <NUM_LIT:0.>))<EOL>def mask_if_invalid(x, mask):<EOL><INDENT>mask = tf.fill(<EOL>tf.shape(input=x), value=np.array(mask, x.dtype.as_numpy_dtype))<EOL>return tf.where(is_valid, x, mask)<EOL><DEDENT>z = (response - mean) / mask_if_invalid(grad_mean, <NUM_LIT:1.>)<EOL>if learning_rate is not None:<EOL><INDENT>z *= learning_rate[..., tf.newaxis]<EOL><DEDENT>z += predicted_linear_response_start<EOL>if offset is not None:<EOL><INDENT>z -= offset<EOL><DEDENT>if dispersion is not None:<EOL><INDENT>variance *= dispersion<EOL><DEDENT>w = (<EOL>mask_if_invalid(grad_mean, <NUM_LIT:0.>) *<EOL>tf.math.rsqrt(mask_if_invalid(variance, np.inf)))<EOL>a = model_matrix * w[..., tf.newaxis]<EOL>b = z * w<EOL>if l2_regularizer is None:<EOL><INDENT>l2_regularizer = np.array(<NUM_LIT:0>, a.dtype.as_numpy_dtype)<EOL><DEDENT>else:<EOL><INDENT>l2_regularizer_ = distribution_util.maybe_get_static_value(<EOL>l2_regularizer, a.dtype.as_numpy_dtype)<EOL>if l2_regularizer_ is not None:<EOL><INDENT>l2_regularizer = l2_regularizer_<EOL><DEDENT><DEDENT>def _embed_l2_regularization():<EOL><INDENT>"""<STR_LIT>"""<EOL>num_model_coefficients = num_cols(model_matrix)<EOL>batch_shape = tf.shape(input=model_matrix)[:-<NUM_LIT:2>]<EOL>eye = tf.eye(<EOL>num_model_coefficients, batch_shape=batch_shape, dtype=a.dtype)<EOL>a_ = tf.concat([a, tf.sqrt(l2_regularizer) * eye], axis=-<NUM_LIT:2>)<EOL>b_ = distribution_util.pad(<EOL>b, count=num_model_coefficients, axis=-<NUM_LIT:1>, back=True)<EOL>l2_regularizer_ = np.array(<NUM_LIT:0>, a.dtype.as_numpy_dtype)<EOL>return a_, b_, l2_regularizer_<EOL><DEDENT>a, b, l2_regularizer = prefer_static.cond(<EOL>prefer_static.reduce_all([not(fast_unsafe_numerics),<EOL>l2_regularizer > <NUM_LIT:0.>]),<EOL>_embed_l2_regularization,<EOL>lambda: (a, b, l2_regularizer))<EOL>model_coefficients_next = tf.linalg.lstsq(<EOL>a,<EOL>b[..., tf.newaxis],<EOL>fast=fast_unsafe_numerics,<EOL>l2_regularizer=l2_regularizer,<EOL>name='<STR_LIT>')<EOL>model_coefficients_next = model_coefficients_next[..., <NUM_LIT:0>]<EOL>predicted_linear_response_next = calculate_linear_predictor(<EOL>model_matrix,<EOL>model_coefficients_next,<EOL>offset,<EOL>name='<STR_LIT>')<EOL>return model_coefficients_next, predicted_linear_response_next<EOL><DEDENT>

Runs one step of Fisher scoring.

    Args:
      model_matrix: (Batch of) `float`-like, matrix-shaped `Tensor` where each row
        represents a sample's features.
      response: (Batch of) vector-shaped `Tensor` where each element represents a
        sample's observed response (to the corresponding row of features). Must
        have same `dtype` as `model_matrix`.
      model: `tfp.glm.ExponentialFamily`-like instance used to construct the
        negative log-likelihood loss, gradient, and expected Hessian (i.e., the
        Fisher information matrix).
      model_coefficients_start: Optional (batch of) vector-shaped `Tensor`
        representing the initial model coefficients, one for each column in
        `model_matrix`. Must have same `dtype` as `model_matrix`.
        Default value: Zeros.
      predicted_linear_response_start: Optional `Tensor` with `shape`, `dtype`
        matching `response`; represents `offset` shifted initial linear
        predictions based on `model_coefficients_start`.
        Default value: `offset` if `model_coefficients is None`, and
        `tf.linalg.matvec(model_matrix, model_coefficients_start) + offset`
        otherwise.
      l2_regularizer: Optional scalar `Tensor` representing L2 regularization
        penalty, i.e.,
        `loss(w) = sum{-log p(y[i]|x[i],w) : i=1..n} + l2_regularizer ||w||_2^2`.
        Default value: `None` (i.e., no L2 regularization).
      dispersion: Optional (batch of) `Tensor` representing `response` dispersion,
        i.e., as in, `p(y|theta) := exp((y theta - A(theta)) / dispersion)`.
        Must broadcast with rows of `model_matrix`.
        Default value: `None` (i.e., "no dispersion").
      offset: Optional `Tensor` representing constant shift applied to
        `predicted_linear_response`.  Must broadcast to `response`.
        Default value: `None` (i.e., `tf.zeros_like(response)`).
      learning_rate: Optional (batch of) scalar `Tensor` used to dampen iterative
        progress. Typically only needed if optimization diverges, should be no
        larger than `1` and typically very close to `1`.
        Default value: `None` (i.e., `1`).
      fast_unsafe_numerics: Optional Python `bool` indicating if solve should be
        based on Cholesky or QR decomposition.
        Default value: `True` (i.e., "prefer speed via Cholesky decomposition").
      name: Python `str` used as name prefix to ops created by this function.
        Default value: `"fit_one_step"`.

    Returns:
      model_coefficients: (Batch of) vector-shaped `Tensor`; represents the
        next estimate of the model coefficients, one for each column in
        `model_matrix`.
      predicted_linear_response: `response`-shaped `Tensor` representing linear
        predictions based on new `model_coefficients`, i.e.,
        `tf.linalg.matvec(model_matrix, model_coefficients_next) + offset`.


def convergence_criteria_fn(<EOL>is_converged_previous,  <EOL>iter_,<EOL>model_coefficients_previous,<EOL>predicted_linear_response_previous,  <EOL>model_coefficients_next,<EOL>predicted_linear_response_next,  <EOL>response,  <EOL>model,  <EOL>dispersion):  <EOL><INDENT>"""<STR_LIT>"""<EOL>relative_euclidean_norm = (<EOL>tf.norm(<EOL>tensor=model_coefficients_previous - model_coefficients_next,<EOL>ord=norm_order,<EOL>axis=-<NUM_LIT:1>) /<EOL>(<NUM_LIT:1.> +<EOL>tf.norm(tensor=model_coefficients_previous, ord=norm_order, axis=-<NUM_LIT:1>)))<EOL>return (iter_ > <NUM_LIT:0>) & tf.reduce_all(<EOL>input_tensor=relative_euclidean_norm < tolerance)<EOL><DEDENT>return convergence_criteria_fn<EOL>

Returns Python `callable` which indicates fitting procedure has converged.

    Writing old, new `model_coefficients` as `w0`, `w1`, this function
    defines convergence as,

    ```python
    relative_euclidean_norm = (tf.norm(w0 - w1, ord=2, axis=-1) /
                               (1. + tf.norm(w0, ord=2, axis=-1)))
    reduce_all(relative_euclidean_norm < tolerance)
    ```

    where `tf.norm(x, ord=2)` denotes the [Euclidean norm](
    https://en.wikipedia.org/wiki/Norm_(mathematics)#Euclidean_norm) of `x`.

    Args:
      tolerance: `float`-like `Tensor` indicating convergence, i.e., when
        max relative Euclidean norm weights difference < tolerance`.
        Default value: `1e-5`.
      norm_order: Order of the norm. Default value: `2` (i.e., "Euclidean norm".)

    Returns:
      convergence_criteria_fn: Python `callable` which returns `bool` `Tensor`
        indicated fitting procedure has converged. (See inner function
        specification for argument signature.)
        Default value: `1e-5`.


graph_deps = [model_matrix, response, model_coefficients,<EOL>predicted_linear_response, offset]<EOL>with tf.compat.v1.name_scope(name, '<STR_LIT>', graph_deps):<EOL><INDENT>dtype = dtype_util.common_dtype(graph_deps, np.float32)<EOL>model_matrix = tf.convert_to_tensor(<EOL>value=model_matrix, dtype=dtype, name='<STR_LIT>')<EOL>if offset is not None:<EOL><INDENT>offset = tf.convert_to_tensor(value=offset, dtype=dtype, name='<STR_LIT>')<EOL><DEDENT>response = tf.convert_to_tensor(<EOL>value=response, dtype=dtype, name='<STR_LIT>')<EOL>use_default_model_coefficients = model_coefficients is None<EOL>if use_default_model_coefficients:<EOL><INDENT>batch_shape = tf.shape(input=model_matrix)[:-<NUM_LIT:2>]<EOL>num_columns = tf.shape(input=model_matrix)[-<NUM_LIT:1>]<EOL>model_coefficients = tf.zeros(<EOL>shape=tf.concat([batch_shape, [num_columns]], axis=<NUM_LIT:0>),<EOL>dtype=dtype, name='<STR_LIT>')<EOL><DEDENT>else:<EOL><INDENT>model_coefficients = tf.convert_to_tensor(<EOL>value=model_coefficients, dtype=dtype, name='<STR_LIT>')<EOL><DEDENT>if predicted_linear_response is None:<EOL><INDENT>if use_default_model_coefficients:<EOL><INDENT>if offset is None:<EOL><INDENT>predicted_linear_response = tf.zeros_like(<EOL>response, dtype, name='<STR_LIT>')<EOL><DEDENT>else:<EOL><INDENT>predicted_linear_response = tf.broadcast_to(<EOL>offset,<EOL>tf.shape(input=response),<EOL>name='<STR_LIT>')<EOL><DEDENT><DEDENT>else:<EOL><INDENT>predicted_linear_response = calculate_linear_predictor(<EOL>model_matrix, model_coefficients, offset)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>predicted_linear_response = tf.convert_to_tensor(<EOL>value=predicted_linear_response,<EOL>dtype=dtype,<EOL>name='<STR_LIT>')<EOL><DEDENT><DEDENT>return [<EOL>model_matrix,<EOL>response,<EOL>model_coefficients,<EOL>predicted_linear_response,<EOL>offset,<EOL>]<EOL>

Helper to `fit` which sanitizes input args.

    Args:
      model_matrix: (Batch of) `float`-like, matrix-shaped `Tensor` where each row
        represents a sample's features.
      response: (Batch of) vector-shaped `Tensor` where each element represents a
        sample's observed response (to the corresponding row of features). Must
        have same `dtype` as `model_matrix`.
      model_coefficients: Optional (batch of) vector-shaped `Tensor` representing
        the model coefficients, one for each column in `model_matrix`. Must have
        same `dtype` as `model_matrix`.
        Default value: `tf.zeros(tf.shape(model_matrix)[-1], model_matrix.dtype)`.
      predicted_linear_response: Optional `Tensor` with `shape`, `dtype` matching
        `response`; represents `offset` shifted initial linear predictions based
        on current `model_coefficients`.
        Default value: `offset` if `model_coefficients is None`, and
        `tf.linalg.matvec(model_matrix, model_coefficients_start) + offset`
        otherwise.
      offset: Optional `Tensor` with `shape`, `dtype` matching `response`;
        represents constant shift applied to `predicted_linear_response`.
        Default value: `None` (i.e., `tf.zeros_like(response)`).
      name: Python `str` used as name prefix to ops created by this function.
        Default value: `"prepare_args"`.

    Returns:
      model_matrix: A `Tensor` with `shape`, `dtype` and values of the
        `model_matrix` argument.
      response: A `Tensor` with `shape`, `dtype` and values of the
        `response` argument.
      model_coefficients_start: A `Tensor` with `shape`, `dtype` and
        values of the `model_coefficients_start` argument if specified.
        A (batch of) vector-shaped `Tensors` with `dtype` matching `model_matrix`
        containing the default starting point otherwise.
      predicted_linear_response:  A `Tensor` with `shape`, `dtype` and
        values of the `predicted_linear_response` argument if specified.
        A `Tensor` with `shape`, `dtype` matching `response` containing the
        default value otherwise.
      offset: A `Tensor` with `shape`, `dtype` and values of the `offset` argument
        if specified or `None` otherwise.


def calculate_linear_predictor(model_matrix, model_coefficients, offset=None,<EOL>name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[model_matrix, model_coefficients, offset]):<EOL><INDENT>predicted_linear_response = tf.linalg.matvec(model_matrix,<EOL>model_coefficients)<EOL>if offset is not None:<EOL><INDENT>predicted_linear_response += offset<EOL><DEDENT>return predicted_linear_response<EOL><DEDENT>

Computes `model_matrix @ model_coefficients + offset`.


if tf.compat.dimension_value(x.shape[-<NUM_LIT:1>]) is not None:<EOL><INDENT>return tf.compat.dimension_value(x.shape[-<NUM_LIT:1>])<EOL><DEDENT>return tf.shape(input=x)[-<NUM_LIT:1>]<EOL>

Returns number of cols in a given `Tensor`.


if not name or name[-<NUM_LIT:1>] != '<STR_LIT:/>':  <EOL><INDENT>with tf.compat.v1.name_scope(name or type(self).__name__) as name:<EOL><INDENT>pass<EOL><DEDENT><DEDENT>self._name = name<EOL>

Creates the ExponentialFamily.

        Args:
          name: Python `str` used as TF namescope for ops created by member
            functions. Default value: `None` (i.e., the subclass name).


return self._is_canonical<EOL>

Returns `True` when `variance(r) == grad_mean(r)` for all `r`.


return self._name<EOL>

Returns TF namescope prefixed to ops created by member functions.


raise NotImplementedError('<STR_LIT>')<EOL>


def __call__(self, predicted_linear_response, name=None):

with self._name_scope(name, '<STR_LIT>', [predicted_linear_response]):<EOL><INDENT>predicted_linear_response = tf.convert_to_tensor(<EOL>value=predicted_linear_response, name='<STR_LIT>')<EOL>return self._call(predicted_linear_response)<EOL><DEDENT>

Computes `mean(r), var(mean), d/dr mean(r)` for linear response, `r`.

        Here `mean` and `var` are the mean and variance of the sufficient statistic,
        which may not be the same as the mean and variance of the random variable
        itself.  If the distribution's density has the form

        ```none
        p_Y(y) = h(y) Exp[dot(theta, T(y)) - A]
        ```

        where `theta` and `A` are constants and `h` and `T` are known functions,
        then `mean` and `var` are the mean and variance of `T(Y)`.  In practice,
        often `T(Y) := Y` and in that case the distinction doesn't matter.

        Args:
          predicted_linear_response: `float`-like `Tensor` corresponding to
            `tf.matmul(model_matrix, weights)`.
          name: Python `str` used as TF namescope for ops created by member
            functions. Default value: `None` (i.e., 'call').

        Returns:
          mean: `Tensor` with shape and dtype of `predicted_linear_response`
            representing the distribution prescribed mean, given the prescribed
            linear-response to mean mapping.
          variance: `Tensor` with shape and dtype of `predicted_linear_response`
            representing the distribution prescribed variance, given the prescribed
            linear-response to mean mapping.
          grad_mean: `Tensor` with shape and dtype of `predicted_linear_response`
            representing the gradient of the mean with respect to the
            linear-response and given the prescribed linear-response to mean
            mapping.


def _log_prob(self, response, predicted_linear_response):

Implements the log-probability computation.


def log_prob(self, response, predicted_linear_response, name=None):

with self._name_scope(<EOL>name, '<STR_LIT>', [response, predicted_linear_response]):<EOL><INDENT>dtype = dtype_util.common_dtype([response, predicted_linear_response])<EOL>response = tf.convert_to_tensor(<EOL>value=response, dtype=dtype, name='<STR_LIT>')<EOL>predicted_linear_response = tf.convert_to_tensor(<EOL>value=predicted_linear_response, name='<STR_LIT>')<EOL>return self._log_prob(response, predicted_linear_response)<EOL><DEDENT>

Computes `D(param=mean(r)).log_prob(response)` for linear response, `r`.

        Args:
          response: `float`-like `Tensor` representing observed ("actual")
            responses.
          predicted_linear_response: `float`-like `Tensor` corresponding to
            `tf.matmul(model_matrix, weights)`.
          name: Python `str` used as TF namescope for ops created by member
            functions. Default value: `None` (i.e., 'log_prob').

        Returns:
          log_prob: `Tensor` with shape and dtype of `predicted_linear_response`
            representing the distribution prescribed log-probability of the observed
            `response`s.


@contextlib.contextmanager<EOL><INDENT>def _name_scope(self, name=None, default_name=None, values=None):<DEDENT>

with tf.compat.v1.name_scope(self.name):<EOL><INDENT>with tf.compat.v1.name_scope(<EOL>name, default_name, values=values or []) as scope:<EOL><INDENT>yield scope<EOL><DEDENT><DEDENT>

Helper function to standardize op scope.


super(CustomExponentialFamily, self).__init__(name)<EOL>self._distribution_fn = distribution_fn<EOL>self._inverse_link_fn = (<EOL>linear_model_to_mean_fn.forward<EOL>if isinstance(linear_model_to_mean_fn, bijectors.Bijector)<EOL>else linear_model_to_mean_fn)<EOL>self._is_canonical = is_canonical<EOL>

Creates the `CustomExponentialFamily`.

        Args:
          distribution_fn: Python `callable` which returns a
            `tf.distribution.Distribution`-like instance from a single input
            representing the distribution's required `mean`, i.e.,
            `mean = linear_model_to_mean_fn(matmul(model_matrix, weights))`.
          linear_model_to_mean_fn: Python `callable` which returns the
            distribution's required mean as computed from the predicted linear
            response, `matmul(model_matrix, weights)`.
          is_canonical: Python `bool` indicating that taken together,
            `distribution_fn` and `linear_model_to_mean_fn` imply that the
            distribution's `variance` is equivalent to `d/dr
            linear_model_to_mean_fn(r)`.
          name: Python `str` used as TF namescope for ops created by member
            functions. Default value: `None` (i.e., the subclass name).


def _fn_wrapped(fn):<EOL><INDENT>"""<STR_LIT>"""<EOL>doc = inspect.cleandoc(fn.__doc__)<EOL>for k, v in six.iteritems(kwargs):<EOL><INDENT>pattern = r'<STR_LIT>' + str(k) + r'<STR_LIT>'<EOL>doc = re.sub(pattern, lambda match: v, doc)  <EOL><DEDENT>fn.__doc__ = doc<EOL>return fn<EOL><DEDENT>return _fn_wrapped<EOL>

Decorator to programmatically expand the docstring.

    Args:
      **kwargs: Keyword arguments to set. For each key-value pair `k` and `v`,
        the key is found as `${k}` in the docstring and replaced with `v`.

    Returns:
      Decorated function.


def softplus_and_shift(x, shift=<NUM_LIT>, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [x, shift]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>y = tf.nn.softplus(x)<EOL>if shift is not None:<EOL><INDENT>y += shift<EOL><DEDENT>return y<EOL><DEDENT>

Converts (batch of) scalars to (batch of) positive valued scalars.

    Args:
      x: (Batch of) `float`-like `Tensor` representing scalars which will be
        transformed into positive elements.
      shift: `Tensor` added to `softplus` transformation of elements.
        Default value: `1e-5`.
      name: A `name_scope` name for operations created by this function.
        Default value: `None` (i.e., "positive_tril_with_shift").

    Returns:
      scale: (Batch of) scalars`with `x.dtype` and `x.shape`.


def tril_with_diag_softplus_and_shift(x, diag_shift=<NUM_LIT>, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[x, diag_shift]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>x = tfd.fill_triangular(x)<EOL>diag = softplus_and_shift(tf.linalg.diag_part(x), diag_shift)<EOL>x = tf.linalg.set_diag(x, diag)<EOL>return x<EOL><DEDENT>

Converts (batch of) vectors to (batch of) lower-triangular scale matrices.

    Args:
      x: (Batch of) `float`-like `Tensor` representing vectors which will be
        transformed into lower-triangular scale matrices with positive diagonal
        elements. Rightmost shape `n` must be such that
        `n = dims * (dims + 1) / 2` for some positive, integer `dims`.
      diag_shift: `Tensor` added to `softplus` transformation of diagonal
        elements.
        Default value: `1e-5`.
      name: A `name_scope` name for operations created by this function.
        Default value: `None` (i.e., "tril_with_diag_softplus_and_shift").

    Returns:
      scale_tril: (Batch of) lower-triangular `Tensor` with `x.dtype` and
        rightmost shape `[dims, dims]` where `n = dims * (dims + 1) / 2` where
        `n = x.shape[-1]`.


with tf.compat.v1.name_scope(name, '<STR_LIT>', [x, dims]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>x = layer_fn(x, dims + dims * (dims + <NUM_LIT:1>) // <NUM_LIT:2>)<EOL>return tfd.MultivariateNormalTriL(<EOL>loc=loc_fn(x[..., :dims]),<EOL>scale_tril=scale_fn(x[..., dims:]))<EOL><DEDENT>

Constructs a trainable `tfd.MultivariateNormalTriL` distribution.

    This function creates a MultivariateNormal (MVN) with lower-triangular scale
    matrix. By default the MVN is parameterized via affine transformation of input
    tensor `x`. Using default args, this function is mathematically equivalent to:

    ```none
    Y = MVN(loc=matmul(W, x) + b,
            scale_tril=f(reshape_tril(matmul(M, x) + c)))

    where,
      W in R^[d, n]
      M in R^[d*(d+1)/2, n]
      b in R^d
      c in R^d
      f(S) = set_diag(S, softplus(matrix_diag_part(S)) + 1e-5)
    ```

    Observe that `f` makes the diagonal of the triangular-lower scale matrix be
    positive and no smaller than `1e-5`.

    #### Examples

    ```python
    # This example fits a multilinear regression loss.
    import tensorflow as tf
    import tensorflow_probability as tfp

    # Create fictitious training data.
    dtype = np.float32
    n = 3000    # number of samples
    x_size = 4  # size of single x
    y_size = 2  # size of single y
    def make_training_data():
      np.random.seed(142)
      x = np.random.randn(n, x_size).astype(dtype)
      w = np.random.randn(x_size, y_size).astype(dtype)
      b = np.random.randn(1, y_size).astype(dtype)
      true_mean = np.tensordot(x, w, axes=[[-1], [0]]) + b
      noise = np.random.randn(n, y_size).astype(dtype)
      y = true_mean + noise
      return y, x
    y, x = make_training_data()

    # Build TF graph for fitting MVNTriL maximum likelihood estimator.
    mvn = tfp.trainable_distributions.multivariate_normal_tril(x, dims=y_size)
    loss = -tf.reduce_mean(mvn.log_prob(y))
    train_op = tf.train.AdamOptimizer(learning_rate=2.**-3).minimize(loss)
    mse = tf.reduce_mean(tf.squared_difference(y, mvn.mean()))
    init_op = tf.global_variables_initializer()

    # Run graph 1000 times.
    num_steps = 1000
    loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
    mse_ = np.zeros(num_steps)
    with tf.Session() as sess:
      sess.run(init_op)
      for it in xrange(loss_.size):
        _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
        if it % 200 == 0 or it == loss_.size - 1:
          print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

    # ==> iteration:0    loss:38.2020797729  mse:4.17175960541
    #     iteration:200  loss:2.90179634094  mse:0.990987896919
    #     iteration:400  loss:2.82727336884  mse:0.990926623344
    #     iteration:600  loss:2.82726788521  mse:0.990926682949
    #     iteration:800  loss:2.82726788521  mse:0.990926682949
    #     iteration:999  loss:2.82726788521  mse:0.990926682949
    ```

    Args:
      x: `Tensor` with floating type. Must have statically defined rank and
        statically known right-most dimension.
      dims: Scalar, `int`, `Tensor` indicated the MVN event size, i.e., the
        created MVN will be distribution over length-`dims` vectors.
      layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
        returns a transformation of `x` with shape
        `tf.concat([tf.shape(x)[:-1], [d]], axis=0)`.
        Default value: `tf.layers.dense`.
      loc_fn: Python `callable` which transforms the `loc` parameter. Takes a
        (batch of) length-`dims` vectors and returns a `Tensor` of same shape and
        `dtype`.
        Default value: `lambda x: x`.
      scale_fn: Python `callable` which transforms the `scale` parameters. Takes a
        (batch of) length-`dims * (dims + 1) / 2` vectors and returns a
        lower-triangular `Tensor` of same batch shape with rightmost dimensions
        having shape `[dims, dims]`.
        Default value: `tril_with_diag_softplus_and_shift`.
      name: A `name_scope` name for operations created by this function.
        Default value: `None` (i.e., "multivariate_normal_tril").

    Returns:
      mvntril: An instance of `tfd.MultivariateNormalTriL`.


def bernoulli(x, layer_fn=tf.compat.v1.layers.dense, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [x]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>logits = tf.squeeze(layer_fn(x, <NUM_LIT:1>), axis=-<NUM_LIT:1>)<EOL>return tfd.Bernoulli(logits=logits)<EOL><DEDENT>

Constructs a trainable `tfd.Bernoulli` distribution.

    This function creates a Bernoulli distribution parameterized by logits.
    Using default args, this function is mathematically equivalent to:

    ```none
    Y = Bernoulli(logits=matmul(W, x) + b)

    where,
      W in R^[d, n]
      b in R^d
    ```

    #### Examples

    This function can be used as a [logistic regression](
    https://en.wikipedia.org/wiki/Logistic_regression) loss.

    ```python
    # This example fits a logistic regression loss.
    import tensorflow as tf
    import tensorflow_probability as tfp

    # Create fictitious training data.
    dtype = np.float32
    n = 3000    # number of samples
    x_size = 4  # size of single x
    def make_training_data():
      np.random.seed(142)
      x = np.random.randn(n, x_size).astype(dtype)
      w = np.random.randn(x_size).astype(dtype)
      b = np.random.randn(1).astype(dtype)
      true_logits = np.tensordot(x, w, axes=[[-1], [-1]]) + b
      noise = np.random.logistic(size=n).astype(dtype)
      y = dtype(true_logits + noise > 0.)
      return y, x
    y, x = make_training_data()

    # Build TF graph for fitting Bernoulli maximum likelihood estimator.
    bernoulli = tfp.trainable_distributions.bernoulli(x)
    loss = -tf.reduce_mean(bernoulli.log_prob(y))
    train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
    mse = tf.reduce_mean(tf.squared_difference(y, bernoulli.mean()))
    init_op = tf.global_variables_initializer()

    # Run graph 1000 times.
    num_steps = 1000
    loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
    mse_ = np.zeros(num_steps)
    with tf.Session() as sess:
      sess.run(init_op)
      for it in xrange(loss_.size):
        _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
        if it % 200 == 0 or it == loss_.size - 1:
          print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

    # ==> iteration:0    loss:0.635675370693  mse:0.222526371479
    #     iteration:200  loss:0.440077394247  mse:0.143687799573
    #     iteration:400  loss:0.440077394247  mse:0.143687844276
    #     iteration:600  loss:0.440077394247  mse:0.143687844276
    #     iteration:800  loss:0.440077424049  mse:0.143687844276
    #     iteration:999  loss:0.440077424049  mse:0.143687844276
    ```

    Args:
      x: `Tensor` with floating type. Must have statically defined rank and
        statically known right-most dimension.
      layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
        returns a transformation of `x` with shape
        `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
        Default value: `tf.layers.dense`.
      name: A `name_scope` name for operations created by this function.
        Default value: `None` (i.e., "bernoulli").

    Returns:
      bernoulli: An instance of `tfd.Bernoulli`.


def normal(x,<EOL>layer_fn=tf.compat.v1.layers.dense,<EOL>loc_fn=lambda x: x,<EOL>scale_fn=<NUM_LIT:1.>,<EOL>name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [x]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>if callable(scale_fn):<EOL><INDENT>y = layer_fn(x, <NUM_LIT:2>)<EOL>loc = loc_fn(y[..., <NUM_LIT:0>])<EOL>scale = scale_fn(y[..., <NUM_LIT:1>])<EOL><DEDENT>else:<EOL><INDENT>y = tf.squeeze(layer_fn(x, <NUM_LIT:1>), axis=-<NUM_LIT:1>)<EOL>loc = loc_fn(y)<EOL>scale = tf.cast(scale_fn, loc.dtype.base_dtype)<EOL><DEDENT>return tfd.Normal(loc=loc, scale=scale)<EOL><DEDENT>

Constructs a trainable `tfd.Normal` distribution.

    This function creates a Normal distribution parameterized by loc and scale.
    Using default args, this function is mathematically equivalent to:

    ```none
    Y = Normal(loc=matmul(W, x) + b, scale=1)

    where,
      W in R^[d, n]
      b in R^d
    ```

    #### Examples

    This function can be used as a [linear regression](
    https://en.wikipedia.org/wiki/Linear_regression) loss.

    ```python
    # This example fits a linear regression loss.
    import tensorflow as tf
    import tensorflow_probability as tfp

    # Create fictitious training data.
    dtype = np.float32
    n = 3000    # number of samples
    x_size = 4  # size of single x
    def make_training_data():
      np.random.seed(142)
      x = np.random.randn(n, x_size).astype(dtype)
      w = np.random.randn(x_size).astype(dtype)
      b = np.random.randn(1).astype(dtype)
      true_mean = np.tensordot(x, w, axes=[[-1], [-1]]) + b
      noise = np.random.randn(n).astype(dtype)
      y = true_mean + noise
      return y, x
    y, x = make_training_data()

    # Build TF graph for fitting Normal maximum likelihood estimator.
    normal = tfp.trainable_distributions.normal(x)
    loss = -tf.reduce_mean(normal.log_prob(y))
    train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
    mse = tf.reduce_mean(tf.squared_difference(y, normal.mean()))
    init_op = tf.global_variables_initializer()

    # Run graph 1000 times.
    num_steps = 1000
    loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
    mse_ = np.zeros(num_steps)
    with tf.Session() as sess:
      sess.run(init_op)
      for it in xrange(loss_.size):
        _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
        if it % 200 == 0 or it == loss_.size - 1:
          print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

    # ==> iteration:0    loss:6.34114170074  mse:10.8444051743
    #     iteration:200  loss:1.40146839619  mse:0.965059816837
    #     iteration:400  loss:1.40052902699  mse:0.963181257248
    #     iteration:600  loss:1.40052902699  mse:0.963181257248
    #     iteration:800  loss:1.40052902699  mse:0.963181257248
    #     iteration:999  loss:1.40052902699  mse:0.963181257248
    ```

    Args:
      x: `Tensor` with floating type. Must have statically defined rank and
        statically known right-most dimension.
      layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
        returns a transformation of `x` with shape
        `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
        Default value: `tf.layers.dense`.
      loc_fn: Python `callable` which transforms the `loc` parameter. Takes a
        (batch of) length-`dims` vectors and returns a `Tensor` of same shape and
        `dtype`.
        Default value: `lambda x: x`.
      scale_fn: Python `callable` or `Tensor`. If a `callable` transforms the
        `scale` parameters; if `Tensor` is the `tfd.Normal` `scale` argument.
        Takes a (batch of) length-`dims` vectors and returns a `Tensor` of same
        size. (Taking a `callable` or `Tensor` is how `tf.Variable` intializers
        behave.)
        Default value: `1`.
      name: A `name_scope` name for operations created by this function.
        Default value: `None` (i.e., "normal").

    Returns:
      normal: An instance of `tfd.Normal`.


def poisson(x,<EOL>layer_fn=tf.compat.v1.layers.dense,<EOL>log_rate_fn=lambda x: x,<EOL>name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [x]):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name='<STR_LIT:x>')<EOL>log_rate = log_rate_fn(tf.squeeze(layer_fn(x, <NUM_LIT:1>), axis=-<NUM_LIT:1>))<EOL>return tfd.Poisson(log_rate=log_rate)<EOL><DEDENT>

Constructs a trainable `tfd.Poisson` distribution.

    This function creates a Poisson distribution parameterized by log rate.
    Using default args, this function is mathematically equivalent to:

    ```none
    Y = Poisson(log_rate=matmul(W, x) + b)

    where,
      W in R^[d, n]
      b in R^d
    ```

    #### Examples

    This can be used as a [Poisson regression](
    https://en.wikipedia.org/wiki/Poisson_regression) loss.

    ```python
    # This example fits a poisson regression loss.
    import numpy as np
    import tensorflow as tf
    import tensorflow_probability as tfp

    # Create fictitious training data.
    dtype = np.float32
    n = 3000    # number of samples
    x_size = 4  # size of single x
    def make_training_data():
      np.random.seed(142)
      x = np.random.randn(n, x_size).astype(dtype)
      w = np.random.randn(x_size).astype(dtype)
      b = np.random.randn(1).astype(dtype)
      true_log_rate = np.tensordot(x, w, axes=[[-1], [-1]]) + b
      y = np.random.poisson(lam=np.exp(true_log_rate)).astype(dtype)
      return y, x
    y, x = make_training_data()

    # Build TF graph for fitting Poisson maximum likelihood estimator.
    poisson = tfp.trainable_distributions.poisson(x)
    loss = -tf.reduce_mean(poisson.log_prob(y))
    train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
    mse = tf.reduce_mean(tf.squared_difference(y, poisson.mean()))
    init_op = tf.global_variables_initializer()

    # Run graph 1000 times.
    num_steps = 1000
    loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
    mse_ = np.zeros(num_steps)
    with tf.Session() as sess:
      sess.run(init_op)
      for it in xrange(loss_.size):
        _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
        if it % 200 == 0 or it == loss_.size - 1:
          print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

    # ==> iteration:0    loss:37.0814208984  mse:6359.41259766
    #     iteration:200  loss:1.42010736465  mse:40.7654914856
    #     iteration:400  loss:1.39027583599  mse:8.77660560608
    #     iteration:600  loss:1.3902695179   mse:8.78443241119
    #     iteration:800  loss:1.39026939869  mse:8.78443622589
    #     iteration:999  loss:1.39026939869  mse:8.78444766998
    ```

    Args:
      x: `Tensor` with floating type. Must have statically defined rank and
        statically known right-most dimension.
      layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
        returns a transformation of `x` with shape
        `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
        Default value: `tf.layers.dense`.
      log_rate_fn: Python `callable` which transforms the `log_rate` parameter.
        Takes a (batch of) length-`dims` vectors and returns a `Tensor` of same
        shape and `dtype`.
        Default value: `lambda x: x`.
      name: A `name_scope` name for operations created by this function.
        Default value: `None` (i.e., "poisson").

    Returns:
      poisson: An instance of `tfd.Poisson`.


def matrix_rank(a, tol=None, validate_args=False, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [a, tol]):<EOL><INDENT>a = tf.convert_to_tensor(value=a, dtype_hint=tf.float32, name='<STR_LIT:a>')<EOL>assertions = _maybe_validate_matrix(a, validate_args)<EOL>if assertions:<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>a = tf.identity(a)<EOL><DEDENT><DEDENT>s = tf.linalg.svd(a, compute_uv=False)<EOL>if tol is None:<EOL><INDENT>if a.shape[-<NUM_LIT:2>:].is_fully_defined():<EOL><INDENT>m = np.max(a.shape[-<NUM_LIT:2>:].as_list())<EOL><DEDENT>else:<EOL><INDENT>m = tf.reduce_max(input_tensor=tf.shape(input=a)[-<NUM_LIT:2>:])<EOL><DEDENT>eps = np.finfo(a.dtype.as_numpy_dtype).eps<EOL>tol = (eps * tf.cast(m, a.dtype) *<EOL>tf.reduce_max(input_tensor=s, axis=-<NUM_LIT:1>, keepdims=True))<EOL><DEDENT>return tf.reduce_sum(input_tensor=tf.cast(s > tol, tf.int32), axis=-<NUM_LIT:1>)<EOL><DEDENT>

Compute the matrix rank; the number of non-zero SVD singular values.

    Arguments:
      a: (Batch of) `float`-like matrix-shaped `Tensor`(s) which are to be
        pseudo-inverted.
      tol: Threshold below which the singular value is counted as "zero".
        Default value: `None` (i.e., `eps * max(rows, cols) * max(singular_val)`).
      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: "matrix_rank".

    Returns:
      matrix_rank: (Batch of) `int32` scalars representing the number of non-zero
        singular values.


def cholesky_concat(chol, cols, name=None):

with tf.compat.v2.name_scope(name or '<STR_LIT>'):<EOL><INDENT>dtype = dtype_util.common_dtype([chol, cols], preferred_dtype=tf.float32)<EOL>chol = tf.convert_to_tensor(value=chol, name='<STR_LIT>', dtype=dtype)<EOL>cols = tf.convert_to_tensor(value=cols, name='<STR_LIT>', dtype=dtype)<EOL>n = prefer_static.shape(chol)[-<NUM_LIT:1>]<EOL>mat_nm, mat_mm = cols[..., :n, :], cols[..., n:, :]<EOL>solved_nm = linear_operator_util.matrix_triangular_solve_with_broadcast(<EOL>chol, mat_nm)<EOL>lower_right_mm = tf.linalg.cholesky(<EOL>mat_mm - tf.matmul(solved_nm, solved_nm, adjoint_a=True))<EOL>lower_left_mn = tf.math.conj(tf.linalg.matrix_transpose(solved_nm))<EOL>out_batch = prefer_static.shape(solved_nm)[:-<NUM_LIT:2>]<EOL>chol = tf.broadcast_to(<EOL>chol,<EOL>tf.concat([out_batch, prefer_static.shape(chol)[-<NUM_LIT:2>:]], axis=<NUM_LIT:0>))<EOL>top_right_zeros_nm = tf.zeros_like(solved_nm)<EOL>return tf.concat([tf.concat([chol, top_right_zeros_nm], axis=-<NUM_LIT:1>),<EOL>tf.concat([lower_left_mn, lower_right_mm], axis=-<NUM_LIT:1>)],<EOL>axis=-<NUM_LIT:2>)<EOL><DEDENT>

Concatenates `chol @ chol.T` with additional rows and columns.

    This operation is conceptually identical to:
    ```python
    def cholesky_concat_slow(chol, cols):  # cols shaped (n + m) x m = z x m
      mat = tf.matmul(chol, chol, adjoint_b=True)  # batch of n x n
      # Concat columns.
      mat = tf.concat([mat, cols[..., :tf.shape(mat)[-2], :]], axis=-1)  # n x z
      # Concat rows.
      mat = tf.concat([mat, tf.linalg.matrix_transpose(cols)], axis=-2)  # z x z
      return tf.linalg.cholesky(mat)
    ```
    but whereas `cholesky_concat_slow` would cost `O(z**3)` work,
    `cholesky_concat` only costs `O(z**2 + m**3)` work.

    The resulting (implicit) matrix must be symmetric and positive definite.
    Thus, the bottom right `m x m` must be self-adjoint, and we do not require a
    separate `rows` argument (which can be inferred from `conj(cols.T)`).

    Args:
      chol: Cholesky decomposition of `mat = chol @ chol.T`.
      cols: The new columns whose first `n` rows we would like concatenated to the
        right of `mat = chol @ chol.T`, and whose conjugate transpose we would
        like concatenated to the bottom of `concat(mat, cols[:n,:])`. A `Tensor`
        with final dims `(n+m, m)`. The first `n` rows are the top right rectangle
        (their conjugate transpose forms the bottom left), and the bottom `m x m`
        is self-adjoint.
      name: Optional name for this op.

    Returns:
      chol_concat: The Cholesky decomposition of:
        ```
        [ [ mat  cols[:n, :] ]
          [   conj(cols.T)   ] ]
        ```


vecs = tf.convert_to_tensor(value=vecs, dtype=tf.int64, name='<STR_LIT>')<EOL>m = tf.convert_to_tensor(value=m, dtype=tf.int64, name='<STR_LIT:m>')<EOL>i = tf.convert_to_tensor(value=i, dtype=tf.int64, name='<STR_LIT:i>')<EOL>trailing_elts = tf.broadcast_to(<EOL>tf.range(m + <NUM_LIT:1>,<EOL>prefer_static.shape(vecs, out_type=tf.int64)[-<NUM_LIT:1>]),<EOL>prefer_static.shape(vecs[..., m + <NUM_LIT:1>:]))<EOL>shp = prefer_static.shape(trailing_elts)<EOL>trailing_elts = tf.where(<EOL>tf.equal(trailing_elts, tf.broadcast_to(i, shp)),<EOL>tf.broadcast_to(tf.gather(vecs, [m], axis=-<NUM_LIT:1>), shp),<EOL>tf.broadcast_to(vecs[..., m + <NUM_LIT:1>:], shp))<EOL>vecs_shape = vecs.shape<EOL>vecs = tf.concat([<EOL>vecs[..., :m],<EOL>tf.gather(vecs, i, batch_dims=prefer_static.rank(vecs) - <NUM_LIT:1>), trailing_elts<EOL>], axis=-<NUM_LIT:1>)<EOL>tensorshape_util.set_shape(vecs, vecs_shape)<EOL>return vecs<EOL>

Swaps `m` and `i` on axis -1. (Helper for pivoted_cholesky.)

    Given a batch of int64 vectors `vecs`, scalar index `m`, and compatibly shaped
    per-vector indices `i`, this function swaps elements `m` and `i` in each
    vector. For the use-case below, these are permutation vectors.

    Args:
      vecs: Vectors on which we perform the swap, int64 `Tensor`.
      m: Scalar int64 `Tensor`, the index into which the `i`th element is going.
      i: Batch int64 `Tensor`, shaped like vecs.shape[:-1] + [1]; the index into
        which the `m`th element is going.

    Returns:
      vecs: The updated vectors.


def pivoted_cholesky(matrix, max_rank, diag_rtol=<NUM_LIT>, name=None):

with tf.compat.v2.name_scope(name or '<STR_LIT>'):<EOL><INDENT>dtype = dtype_util.common_dtype([matrix, diag_rtol],<EOL>preferred_dtype=tf.float32)<EOL>matrix = tf.convert_to_tensor(value=matrix, name='<STR_LIT>', dtype=dtype)<EOL>if tensorshape_util.rank(matrix.shape) is None:<EOL><INDENT>raise NotImplementedError('<STR_LIT>')<EOL><DEDENT>max_rank = tf.convert_to_tensor(<EOL>value=max_rank, name='<STR_LIT>', dtype=tf.int64)<EOL>max_rank = tf.minimum(max_rank,<EOL>prefer_static.shape(matrix, out_type=tf.int64)[-<NUM_LIT:1>])<EOL>diag_rtol = tf.convert_to_tensor(<EOL>value=diag_rtol, dtype=dtype, name='<STR_LIT>')<EOL>matrix_diag = tf.linalg.diag_part(matrix)<EOL>orig_error = tf.reduce_max(input_tensor=matrix_diag, axis=-<NUM_LIT:1>)<EOL>def cond(m, pchol, perm, matrix_diag):<EOL><INDENT>"""<STR_LIT>"""<EOL>del pchol<EOL>del perm<EOL>error = tf.linalg.norm(tensor=matrix_diag, ord=<NUM_LIT:1>, axis=-<NUM_LIT:1>)<EOL>max_err = tf.reduce_max(input_tensor=error / orig_error)<EOL>return (m < max_rank) & (tf.equal(m, <NUM_LIT:0>) | (max_err > diag_rtol))<EOL><DEDENT>batch_dims = tensorshape_util.rank(matrix.shape) - <NUM_LIT:2><EOL>def batch_gather(params, indices, axis=-<NUM_LIT:1>):<EOL><INDENT>return tf.gather(params, indices, axis=axis, batch_dims=batch_dims)<EOL><DEDENT>def body(m, pchol, perm, matrix_diag):<EOL><INDENT>"""<STR_LIT>"""<EOL>permuted_diag = batch_gather(matrix_diag, perm[..., m:])<EOL>maxi = tf.argmax(<EOL>input=permuted_diag, axis=-<NUM_LIT:1>, output_type=tf.int64)[..., tf.newaxis]<EOL>maxval = batch_gather(permuted_diag, maxi)<EOL>maxi = maxi + m<EOL>maxval = maxval[..., <NUM_LIT:0>]<EOL>perm = _swap_m_with_i(perm, m, maxi)<EOL>row = batch_gather(matrix, perm[..., m:m + <NUM_LIT:1>], axis=-<NUM_LIT:2>)<EOL>row = batch_gather(row, perm[..., m + <NUM_LIT:1>:])<EOL>prev_rows = pchol[..., :m, :]<EOL>prev_rows_perm_m_onward = batch_gather(prev_rows, perm[..., m + <NUM_LIT:1>:])<EOL>prev_rows_pivot_col = batch_gather(prev_rows, perm[..., m:m + <NUM_LIT:1>])<EOL>row -= tf.reduce_sum(<EOL>input_tensor=prev_rows_perm_m_onward * prev_rows_pivot_col,<EOL>axis=-<NUM_LIT:2>)[..., tf.newaxis, :]<EOL>pivot = tf.sqrt(maxval)[..., tf.newaxis, tf.newaxis]<EOL>row = tf.concat([pivot, row / pivot], axis=-<NUM_LIT:1>)<EOL>paddings = tf.concat([<EOL>tf.zeros([prefer_static.rank(pchol) - <NUM_LIT:1>, <NUM_LIT:2>], dtype=tf.int32),<EOL>[[tf.cast(m, tf.int32), <NUM_LIT:0>]]], axis=<NUM_LIT:0>)<EOL>diag_update = tf.pad(tensor=row**<NUM_LIT:2>, paddings=paddings)[..., <NUM_LIT:0>, :]<EOL>reverse_perm = _invert_permutation(perm)<EOL>matrix_diag -= batch_gather(diag_update, reverse_perm)<EOL>row = tf.pad(tensor=row, paddings=paddings)<EOL>row = batch_gather(row, reverse_perm)<EOL>pchol_shape = pchol.shape<EOL>pchol = tf.concat([pchol[..., :m, :], row, pchol[..., m + <NUM_LIT:1>:, :]],<EOL>axis=-<NUM_LIT:2>)<EOL>tensorshape_util.set_shape(pchol, pchol_shape)<EOL>return m + <NUM_LIT:1>, pchol, perm, matrix_diag<EOL><DEDENT>m = np.int64(<NUM_LIT:0>)<EOL>pchol = tf.zeros_like(matrix[..., :max_rank, :])<EOL>matrix_shape = prefer_static.shape(matrix, out_type=tf.int64)<EOL>perm = tf.broadcast_to(<EOL>prefer_static.range(matrix_shape[-<NUM_LIT:1>]), matrix_shape[:-<NUM_LIT:1>])<EOL>_, pchol, _, _ = tf.while_loop(<EOL>cond=cond, body=body, loop_vars=(m, pchol, perm, matrix_diag))<EOL>pchol = tf.linalg.matrix_transpose(pchol)<EOL>tensorshape_util.set_shape(<EOL>pchol, tensorshape_util.concatenate(matrix_diag.shape, [None]))<EOL>return pchol<EOL><DEDENT>

Computes the (partial) pivoted cholesky decomposition of `matrix`.

    The pivoted Cholesky is a low rank approximation of the Cholesky decomposition
    of `matrix`, i.e. as described in [(Harbrecht et al., 2012)][1]. The
    currently-worst-approximated diagonal element is selected as the pivot at each
    iteration. This yields from a `[B1...Bn, N, N]` shaped `matrix` a `[B1...Bn,
    N, K]` shaped rank-`K` approximation `lr` such that `lr @ lr.T ~= matrix`.
    Note that, unlike the Cholesky decomposition, `lr` is not triangular even in
    a rectangular-matrix sense. However, under a permutation it could be made
    triangular (it has one more zero in each column as you move to the right).

    Such a matrix can be useful as a preconditioner for conjugate gradient
    optimization, i.e. as in [(Wang et al. 2019)][2], as matmuls and solves can be
    cheaply done via the Woodbury matrix identity, as implemented by
    `tf.linalg.LinearOperatorLowRankUpdate`.

    Args:
      matrix: Floating point `Tensor` batch of symmetric, positive definite
        matrices.
      max_rank: Scalar `int` `Tensor`, the rank at which to truncate the
        approximation.
      diag_rtol: Scalar floating point `Tensor` (same dtype as `matrix`). If the
        errors of all diagonal elements of `lr @ lr.T` are each lower than
        `element * diag_rtol`, iteration is permitted to terminate early.
      name: Optional name for the op.

    Returns:
      lr: Low rank pivoted Cholesky approximation of `matrix`.

    #### References

    [1]: H Harbrecht, M Peters, R Schneider. On the low-rank approximation by the
         pivoted Cholesky decomposition. _Applied numerical mathematics_,
         62(4):428-440, 2012.

    [2]: K. A. Wang et al. Exact Gaussian Processes on a Million Data Points.
         _arXiv preprint arXiv:1903.08114_, 2019. https://arxiv.org/abs/1903.08114


def pinv(a, rcond=None, validate_args=False, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [a, rcond]):<EOL><INDENT>a = tf.convert_to_tensor(value=a, name='<STR_LIT:a>')<EOL>assertions = _maybe_validate_matrix(a, validate_args)<EOL>if assertions:<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>a = tf.identity(a)<EOL><DEDENT><DEDENT>dtype = a.dtype.as_numpy_dtype<EOL>if rcond is None:<EOL><INDENT>def get_dim_size(dim):<EOL><INDENT>if tf.compat.dimension_value(a.shape[dim]) is not None:<EOL><INDENT>return tf.compat.dimension_value(a.shape[dim])<EOL><DEDENT>return tf.shape(input=a)[dim]<EOL><DEDENT>num_rows = get_dim_size(-<NUM_LIT:2>)<EOL>num_cols = get_dim_size(-<NUM_LIT:1>)<EOL>if isinstance(num_rows, int) and isinstance(num_cols, int):<EOL><INDENT>max_rows_cols = float(max(num_rows, num_cols))<EOL><DEDENT>else:<EOL><INDENT>max_rows_cols = tf.cast(tf.maximum(num_rows, num_cols), dtype)<EOL><DEDENT>rcond = <NUM_LIT> * max_rows_cols * np.finfo(dtype).eps<EOL><DEDENT>rcond = tf.convert_to_tensor(value=rcond, dtype=dtype, name='<STR_LIT>')<EOL>[<EOL>singular_values,         <EOL>left_singular_vectors,   <EOL>right_singular_vectors,  <EOL>] = tf.linalg.svd(a, full_matrices=False, compute_uv=True)<EOL>cutoff = rcond * tf.reduce_max(input_tensor=singular_values, axis=-<NUM_LIT:1>)<EOL>singular_values = tf.where(<EOL>singular_values > cutoff[..., tf.newaxis], singular_values,<EOL>tf.fill(tf.shape(input=singular_values), np.array(np.inf, dtype)))<EOL>a_pinv = tf.matmul(<EOL>right_singular_vectors / singular_values[..., tf.newaxis, :],<EOL>left_singular_vectors,<EOL>adjoint_b=True)<EOL>if a.shape.ndims is not None:<EOL><INDENT>a_pinv.set_shape(a.shape[:-<NUM_LIT:2>].concatenate([a.shape[-<NUM_LIT:1>], a.shape[-<NUM_LIT:2>]]))<EOL><DEDENT>return a_pinv<EOL><DEDENT>

Compute the Moore-Penrose pseudo-inverse of a matrix.

    Calculate the [generalized inverse of a matrix](
    https://en.wikipedia.org/wiki/Moore%E2%80%93Penrose_inverse) using its
    singular-value decomposition (SVD) and including all large singular values.

    The pseudo-inverse of a matrix `A`, is defined as: "the matrix that 'solves'
    [the least-squares problem] `A @ x = b`," i.e., if `x_hat` is a solution, then
    `A_pinv` is the matrix such that `x_hat = A_pinv @ b`. It can be shown that if
    `U @ Sigma @ V.T = A` is the singular value decomposition of `A`, then
    `A_pinv = V @ inv(Sigma) U^T`. [(Strang, 1980)][1]

    This function is analogous to [`numpy.linalg.pinv`](
    https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.pinv.html).
    It differs only in default value of `rcond`. In `numpy.linalg.pinv`, the
    default `rcond` is `1e-15`. Here the default is
    `10. * max(num_rows, num_cols) * np.finfo(dtype).eps`.

    Args:
      a: (Batch of) `float`-like matrix-shaped `Tensor`(s) which are to be
        pseudo-inverted.
      rcond: `Tensor` of small singular value cutoffs.  Singular values smaller
        (in modulus) than `rcond` * largest_singular_value (again, in modulus) are
        set to zero. Must broadcast against `tf.shape(a)[:-2]`.
        Default value: `10. * max(num_rows, num_cols) * np.finfo(a.dtype).eps`.
      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: "pinv".

    Returns:
      a_pinv: The pseudo-inverse of input `a`. Has same shape as `a` except
        rightmost two dimensions are transposed.

    Raises:
      TypeError: if input `a` does not have `float`-like `dtype`.
      ValueError: if input `a` has fewer than 2 dimensions.

    #### Examples

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

    a = tf.constant([[1.,  0.4,  0.5],
                     [0.4, 0.2,  0.25],
                     [0.5, 0.25, 0.35]])
    tf.matmul(tfp.math.pinv(a), a)
    # ==> array([[1., 0., 0.],
                 [0., 1., 0.],
                 [0., 0., 1.]], dtype=float32)

    a = tf.constant([[1.,  0.4,  0.5,  1.],
                     [0.4, 0.2,  0.25, 2.],
                     [0.5, 0.25, 0.35, 3.]])
    tf.matmul(tfp.math.pinv(a), a)
    # ==> array([[ 0.76,  0.37,  0.21, -0.02],
                 [ 0.37,  0.43, -0.33,  0.02],
                 [ 0.21, -0.33,  0.81,  0.01],
                 [-0.02,  0.02,  0.01,  1.  ]], dtype=float32)
    ```

    #### References

    [1]: G. Strang. "Linear Algebra and Its Applications, 2nd Ed." Academic Press,
         Inc., 1980, pp. 139-142.


def lu_solve(lower_upper, perm, rhs,<EOL>validate_args=False,<EOL>name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [lower_upper, perm, rhs]):<EOL><INDENT>lower_upper = tf.convert_to_tensor(<EOL>value=lower_upper, dtype_hint=tf.float32, name='<STR_LIT>')<EOL>perm = tf.convert_to_tensor(value=perm, dtype_hint=tf.int32, name='<STR_LIT>')<EOL>rhs = tf.convert_to_tensor(<EOL>value=rhs, dtype_hint=lower_upper.dtype, name='<STR_LIT>')<EOL>assertions = _lu_solve_assertions(lower_upper, perm, rhs, validate_args)<EOL>if assertions:<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>lower_upper = tf.identity(lower_upper)<EOL>perm = tf.identity(perm)<EOL>rhs = tf.identity(rhs)<EOL><DEDENT><DEDENT>if rhs.shape.ndims == <NUM_LIT:2> and perm.shape.ndims == <NUM_LIT:1>:<EOL><INDENT>permuted_rhs = tf.gather(rhs, perm, axis=-<NUM_LIT:2>)<EOL><DEDENT>else:<EOL><INDENT>rhs_shape = tf.shape(input=rhs)<EOL>broadcast_batch_shape = tf.broadcast_dynamic_shape(<EOL>rhs_shape[:-<NUM_LIT:2>],<EOL>tf.shape(input=perm)[:-<NUM_LIT:1>])<EOL>d, m = rhs_shape[-<NUM_LIT:2>], rhs_shape[-<NUM_LIT:1>]<EOL>rhs_broadcast_shape = tf.concat([broadcast_batch_shape, [d, m]], axis=<NUM_LIT:0>)<EOL>broadcast_rhs = tf.broadcast_to(rhs, rhs_broadcast_shape)<EOL>broadcast_rhs = tf.reshape(broadcast_rhs, [-<NUM_LIT:1>, d, m])<EOL>broadcast_perm = tf.broadcast_to(perm, rhs_broadcast_shape[:-<NUM_LIT:1>])<EOL>broadcast_perm = tf.reshape(broadcast_perm, [-<NUM_LIT:1>, d])<EOL>broadcast_batch_size = tf.reduce_prod(input_tensor=broadcast_batch_shape)<EOL>broadcast_batch_indices = tf.broadcast_to(<EOL>tf.range(broadcast_batch_size)[:, tf.newaxis],<EOL>[broadcast_batch_size, d])<EOL>broadcast_perm = tf.stack([broadcast_batch_indices, broadcast_perm],<EOL>axis=-<NUM_LIT:1>)<EOL>permuted_rhs = tf.gather_nd(broadcast_rhs, broadcast_perm)<EOL>permuted_rhs = tf.reshape(permuted_rhs, rhs_broadcast_shape)<EOL><DEDENT>lower = tf.linalg.set_diag(<EOL>tf.linalg.band_part(lower_upper, num_lower=-<NUM_LIT:1>, num_upper=<NUM_LIT:0>),<EOL>tf.ones(tf.shape(input=lower_upper)[:-<NUM_LIT:1>], dtype=lower_upper.dtype))<EOL>return linear_operator_util.matrix_triangular_solve_with_broadcast(<EOL>lower_upper,  <EOL>linear_operator_util.matrix_triangular_solve_with_broadcast(<EOL>lower, permuted_rhs),<EOL>lower=False)<EOL><DEDENT>

Solves systems of linear eqns `A X = RHS`, given LU factorizations.

    Note: this function does not verify the implied matrix is actually invertible
    nor is this condition checked even when `validate_args=True`.

    Args:
      lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if
        `matmul(P, matmul(L, U)) = X` then `lower_upper = L + U - eye`.
      perm: `p` as returned by `tf.linag.lu`, i.e., if
        `matmul(P, matmul(L, U)) = X` then `perm = argmax(P)`.
      rhs: Matrix-shaped float `Tensor` representing targets for which to solve;
        `A X = RHS`. To handle vector cases, use:
        `lu_solve(..., rhs[..., tf.newaxis])[..., 0]`.
      validate_args: Python `bool` indicating whether arguments should be checked
        for correctness. Note: this function does not verify the implied matrix is
        actually invertible, even when `validate_args=True`.
        Default value: `False` (i.e., don't validate arguments).
      name: Python `str` name given to ops managed by this object.
        Default value: `None` (i.e., "lu_solve").

    Returns:
      x: The `X` in `A @ X = RHS`.

    #### Examples

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

    x = [[[1., 2],
          [3, 4]],
         [[7, 8],
          [3, 4]]]
    inv_x = tfp.math.lu_solve(*tf.linalg.lu(x), rhs=tf.eye(2))
    tf.assert_near(tf.matrix_inverse(x), inv_x)
    # ==> True
    ```


def lu_matrix_inverse(lower_upper, perm, validate_args=False, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [lower_upper, perm]):<EOL><INDENT>lower_upper = tf.convert_to_tensor(<EOL>value=lower_upper, dtype_hint=tf.float32, name='<STR_LIT>')<EOL>perm = tf.convert_to_tensor(value=perm, dtype_hint=tf.int32, name='<STR_LIT>')<EOL>assertions = _lu_reconstruct_assertions(lower_upper, perm, validate_args)<EOL>if assertions:<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>lower_upper = tf.identity(lower_upper)<EOL>perm = tf.identity(perm)<EOL><DEDENT><DEDENT>shape = tf.shape(input=lower_upper)<EOL>return lu_solve(<EOL>lower_upper, perm,<EOL>rhs=tf.eye(shape[-<NUM_LIT:1>], batch_shape=shape[:-<NUM_LIT:2>], dtype=lower_upper.dtype),<EOL>validate_args=False)<EOL><DEDENT>

Computes a matrix inverse given the matrix's LU decomposition.

    This op is conceptually identical to,

    ````python
    inv_X = tf.lu_matrix_inverse(*tf.linalg.lu(X))
    tf.assert_near(tf.matrix_inverse(X), inv_X)
    # ==> True
    ```

    Note: this function does not verify the implied matrix is actually invertible
    nor is this condition checked even when `validate_args=True`.

    Args:
      lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if
        `matmul(P, matmul(L, U)) = X` then `lower_upper = L + U - eye`.
      perm: `p` as returned by `tf.linag.lu`, i.e., if
        `matmul(P, matmul(L, U)) = X` then `perm = argmax(P)`.
      validate_args: Python `bool` indicating whether arguments should be checked
        for correctness. Note: this function does not verify the implied matrix is
        actually invertible, even when `validate_args=True`.
        Default value: `False` (i.e., don't validate arguments).
      name: Python `str` name given to ops managed by this object.
        Default value: `None` (i.e., "lu_matrix_inverse").

    Returns:
      inv_x: The matrix_inv, i.e.,
        `tf.matrix_inverse(tfp.math.lu_reconstruct(lu, perm))`.

    #### Examples

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

    x = [[[3., 4], [1, 2]],
         [[7., 8], [3, 4]]]
    inv_x = tfp.math.lu_matrix_inverse(*tf.linalg.lu(x))
    tf.assert_near(tf.matrix_inverse(x), inv_x)
    # ==> True
    ```


def lu_reconstruct(lower_upper, perm, validate_args=False, name=None):

with tf.compat.v1.name_scope(name, '<STR_LIT>', [lower_upper, perm]):<EOL><INDENT>lower_upper = tf.convert_to_tensor(<EOL>value=lower_upper, dtype_hint=tf.float32, name='<STR_LIT>')<EOL>perm = tf.convert_to_tensor(value=perm, dtype_hint=tf.int32, name='<STR_LIT>')<EOL>assertions = _lu_reconstruct_assertions(lower_upper, perm, validate_args)<EOL>if assertions:<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>lower_upper = tf.identity(lower_upper)<EOL>perm = tf.identity(perm)<EOL><DEDENT><DEDENT>shape = tf.shape(input=lower_upper)<EOL>lower = tf.linalg.set_diag(<EOL>tf.linalg.band_part(lower_upper, num_lower=-<NUM_LIT:1>, num_upper=<NUM_LIT:0>),<EOL>tf.ones(shape[:-<NUM_LIT:1>], dtype=lower_upper.dtype))<EOL>upper = tf.linalg.band_part(lower_upper, num_lower=<NUM_LIT:0>, num_upper=-<NUM_LIT:1>)<EOL>x = tf.matmul(lower, upper)<EOL>if lower_upper.shape.ndims is None or lower_upper.shape.ndims != <NUM_LIT:2>:<EOL><INDENT>batch_size = tf.reduce_prod(input_tensor=shape[:-<NUM_LIT:2>])<EOL>d = shape[-<NUM_LIT:1>]<EOL>x = tf.reshape(x, [batch_size, d, d])<EOL>perm = tf.reshape(perm, [batch_size, d])<EOL>perm = tf.map_fn(tf.math.invert_permutation, perm)<EOL>batch_indices = tf.broadcast_to(<EOL>tf.range(batch_size)[:, tf.newaxis],<EOL>[batch_size, d])<EOL>x = tf.gather_nd(x, tf.stack([batch_indices, perm], axis=-<NUM_LIT:1>))<EOL>x = tf.reshape(x, shape)<EOL><DEDENT>else:<EOL><INDENT>x = tf.gather(x, tf.math.invert_permutation(perm))<EOL><DEDENT>x.set_shape(lower_upper.shape)<EOL>return x<EOL><DEDENT>

The inverse LU decomposition, `X == lu_reconstruct(*tf.linalg.lu(X))`.

    Args:
      lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if
        `matmul(P, matmul(L, U)) = X` then `lower_upper = L + U - eye`.
      perm: `p` as returned by `tf.linag.lu`, i.e., if
        `matmul(P, matmul(L, U)) = X` then `perm = argmax(P)`.
      validate_args: Python `bool` indicating whether arguments should be checked
        for correctness.
        Default value: `False` (i.e., don't validate arguments).
      name: Python `str` name given to ops managed by this object.
        Default value: `None` (i.e., "lu_reconstruct").

    Returns:
      x: The original input to `tf.linalg.lu`, i.e., `x` as in,
        `lu_reconstruct(*tf.linalg.lu(x))`.

    #### Examples

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

    x = [[[3., 4], [1, 2]],
         [[7., 8], [3, 4]]]
    x_reconstructed = tfp.math.lu_reconstruct(*tf.linalg.lu(x))
    tf.assert_near(x, x_reconstructed)
    # ==> True
    ```


def _lu_reconstruct_assertions(lower_upper, perm, validate_args):

assertions = []<EOL>message = '<STR_LIT>'<EOL>if lower_upper.shape.ndims is not None:<EOL><INDENT>if lower_upper.shape.ndims < <NUM_LIT:2>:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>tf.compat.v1.assert_rank_at_least(lower_upper, rank=<NUM_LIT:2>, message=message))<EOL><DEDENT>message = '<STR_LIT>'<EOL>if lower_upper.shape.ndims is not None and perm.shape.ndims is not None:<EOL><INDENT>if lower_upper.shape.ndims != perm.shape.ndims + <NUM_LIT:1>:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>tf.compat.v1.assert_rank(<EOL>lower_upper, rank=tf.rank(perm) + <NUM_LIT:1>, message=message))<EOL><DEDENT>message = '<STR_LIT>'<EOL>if lower_upper.shape[:-<NUM_LIT:2>].is_fully_defined():<EOL><INDENT>if lower_upper.shape[-<NUM_LIT:2>] != lower_upper.shape[-<NUM_LIT:1>]:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>m, n = tf.split(tf.shape(input=lower_upper)[-<NUM_LIT:2>:], num_or_size_splits=<NUM_LIT:2>)<EOL>assertions.append(tf.compat.v1.assert_equal(m, n, message=message))<EOL><DEDENT>return assertions<EOL>

Returns list of assertions related to `lu_reconstruct` assumptions.


def _lu_solve_assertions(lower_upper, perm, rhs, validate_args):

assertions = _lu_reconstruct_assertions(lower_upper, perm, validate_args)<EOL>message = '<STR_LIT>'<EOL>if rhs.shape.ndims is not None:<EOL><INDENT>if rhs.shape.ndims < <NUM_LIT:2>:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>tf.compat.v1.assert_rank_at_least(rhs, rank=<NUM_LIT:2>, message=message))<EOL><DEDENT>message = '<STR_LIT>'<EOL>if (tf.compat.dimension_value(lower_upper.shape[-<NUM_LIT:1>]) is not None and<EOL>tf.compat.dimension_value(rhs.shape[-<NUM_LIT:2>]) is not None):<EOL><INDENT>if lower_upper.shape[-<NUM_LIT:1>] != rhs.shape[-<NUM_LIT:2>]:<EOL><INDENT>raise ValueError(message)<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(<EOL>tf.compat.v1.assert_equal(<EOL>tf.shape(input=lower_upper)[-<NUM_LIT:1>],<EOL>tf.shape(input=rhs)[-<NUM_LIT:2>],<EOL>message=message))<EOL><DEDENT>return assertions<EOL>

Returns list of assertions related to `lu_solve` assumptions.


with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[sparse_or_dense_a, dense_b]):<EOL><INDENT>dense_b = tf.convert_to_tensor(<EOL>value=dense_b, dtype_hint=tf.float32, name='<STR_LIT>')<EOL>if validate_args:<EOL><INDENT>assert_a_rank_at_least_2 = tf.compat.v1.assert_rank_at_least(<EOL>sparse_or_dense_a,<EOL>rank=<NUM_LIT:2>,<EOL>message='<STR_LIT>')<EOL>assert_b_rank_at_least_2 = tf.compat.v1.assert_rank_at_least(<EOL>dense_b,<EOL>rank=<NUM_LIT:2>,<EOL>message='<STR_LIT>')<EOL>with tf.control_dependencies(<EOL>[assert_a_rank_at_least_2, assert_b_rank_at_least_2]):<EOL><INDENT>sparse_or_dense_a = tf.identity(sparse_or_dense_a)<EOL>dense_b = tf.identity(dense_b)<EOL><DEDENT><DEDENT>if isinstance(sparse_or_dense_a,<EOL>(tf.SparseTensor, tf.compat.v1.SparseTensorValue)):<EOL><INDENT>return _sparse_tensor_dense_matmul(sparse_or_dense_a, dense_b, **kwargs)<EOL><DEDENT>else:<EOL><INDENT>return tf.matmul(sparse_or_dense_a, dense_b, **kwargs)<EOL><DEDENT><DEDENT>

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

    Args:
      sparse_or_dense_a: `SparseTensor` or `Tensor` representing a (batch of)
        matrices.
      dense_b: `Tensor` representing a (batch of) matrices, with the same batch
        shape as `sparse_or_dense_a`. The shape must be compatible with the shape
        of `sparse_or_dense_a` 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_matmul".
      **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul` or
        `tf.matmul`.

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