Instruction:
Can you tell what is the following Python 3 function doing
def _ndtri(p):

  # Constants used in piece-wise rational approximations. Taken from the cephes
  # library:
  # https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
  p0 = list(reversed([-5.99633501014107895267E1,
                      9.80010754185999661536E1,
                      -5.66762857469070293439E1,
                      1.39312609387279679503E1,
                      -1.23916583867381258016E0]))
  q0 = list(reversed([1.0,
                      1.95448858338141759834E0,
                      4.67627912898881538453E0,
                      8.63602421390890590575E1,
                      -2.25462687854119370527E2,
                      2.00260212380060660359E2,
                      -8.20372256168333339912E1,
                      1.59056225126211695515E1,
                      -1.18331621121330003142E0]))
  p1 = list(reversed([4.05544892305962419923E0,
                      3.15251094599893866154E1,
                      5.71628192246421288162E1,
                      4.40805073893200834700E1,
                      1.46849561928858024014E1,
                      2.18663306850790267539E0,
                      -1.40256079171354495875E-1,
                      -3.50424626827848203418E-2,
                      -8.57456785154685413611E-4]))
  q1 = list(reversed([1.0,
                      1.57799883256466749731E1,
                      4.53907635128879210584E1,
                      4.13172038254672030440E1,
                      1.50425385692907503408E1,
                      2.50464946208309415979E0,
                      -1.42182922854787788574E-1,
                      -3.80806407691578277194E-2,
                      -9.33259480895457427372E-4]))
  p2 = list(reversed([3.23774891776946035970E0,
                      6.91522889068984211695E0,
                      3.93881025292474443415E0,
                      1.33303460815807542389E0,
                      2.01485389549179081538E-1,
                      1.23716634817820021358E-2,
                      3.01581553508235416007E-4,
                      2.65806974686737550832E-6,
                      6.23974539184983293730E-9]))
  q2 = list(reversed([1.0,
                      6.02427039364742014255E0,
                      3.67983563856160859403E0,
                      1.37702099489081330271E0,
                      2.16236993594496635890E-1,
                      1.34204006088543189037E-2,
                      3.28014464682127739104E-4,
                      2.89247864745380683936E-6,
                      6.79019408009981274425E-9]))

  def _create_polynomial(var, coeffs):
    """Compute n_th order polynomial via Horner's method."""
    coeffs = np.array(coeffs, dtype_util.as_numpy_dtype(var.dtype))
    if not coeffs.size:
      return tf.zeros_like(var)
    return coeffs[0] + _create_polynomial(var, coeffs[1:]) * var

  maybe_complement_p = tf.where(p > -np.expm1(-2.), 1. - p, p)
  # Write in an arbitrary value in place of 0 for p since 0 will cause NaNs
  # later on. The result from the computation when p == 0 is not used so any
  # number that doesn't result in NaNs is fine.
  sanitized_mcp = tf.where(
      maybe_complement_p <= 0.,
      tf.fill(tf.shape(input=p),
              dtype_util.as_numpy_dtype(p.dtype)(0.5)), maybe_complement_p)

  # Compute x for p > exp(-2): x/sqrt(2pi) = w + w**3 P0(w**2)/Q0(w**2).
  w = sanitized_mcp - 0.5
  ww = w ** 2
  x_for_big_p = w + w * ww * (_create_polynomial(ww, p0)
                              / _create_polynomial(ww, q0))
  x_for_big_p *= -np.sqrt(2. * np.pi)

  # Compute x for p <= exp(-2): x = z - log(z)/z - (1/z) P(1/z) / Q(1/z),
  # where z = sqrt(-2. * log(p)), and P/Q are chosen between two different
  # arrays based on whether p < exp(-32).
  z = tf.sqrt(-2. * tf.math.log(sanitized_mcp))
  first_term = z - tf.math.log(z) / z
  second_term_small_p = (
      _create_polynomial(1. / z, p2) /
      _create_polynomial(1. / z, q2) / z)
  second_term_otherwise = (
      _create_polynomial(1. / z, p1) /
      _create_polynomial(1. / z, q1) / z)
  x_for_small_p = first_term - second_term_small_p
  x_otherwise = first_term - second_term_otherwise

  x = tf.where(sanitized_mcp > np.exp(-2.), x_for_big_p,
               tf.where(z >= 8.0, x_for_small_p, x_otherwise))

  x = tf.where(p > 1. - np.exp(-2.), x, -x)
  infinity_scalar = tf.constant(np.inf, dtype=p.dtype)
  infinity = tf.fill(tf.shape(input=p), infinity_scalar)
  x_nan_replaced = tf.where(
      p <= 0.0, -infinity, tf.where(p >= 1.0, infinity, x))
  return x_nan_replaced

Response: Implements ndtri core logic.


Instruction:
Write a Python 3 script to
log Normal distribution function.

Response:
def log_ndtr(x, series_order=3, name="log_ndtr"):
  """Log Normal distribution function.

  For details of the Normal distribution function see `ndtr`.

  This function calculates `(log o ndtr)(x)` by either calling `log(ndtr(x))` or
  using an asymptotic series. Specifically:
  - For `x > upper_segment`, use the approximation `-ndtr(-x)` based on
    `log(1-x) ~= -x, x << 1`.
  - For `lower_segment < x <= upper_segment`, use the existing `ndtr` technique
    and take a log.
  - For `x <= lower_segment`, we use the series approximation of erf to compute
    the log CDF directly.

  The `lower_segment` is set based on the precision of the input:

  ```
  lower_segment = { -20,  x.dtype=float64
                  { -10,  x.dtype=float32
  upper_segment = {   8,  x.dtype=float64
                  {   5,  x.dtype=float32
  ```

  When `x < lower_segment`, the `ndtr` asymptotic series approximation is:

  ```
     ndtr(x) = scale * (1 + sum) + R_N
     scale   = exp(-0.5 x**2) / (-x sqrt(2 pi))
     sum     = Sum{(-1)^n (2n-1)!! / (x**2)^n, n=1:N}
     R_N     = O(exp(-0.5 x**2) (2N+1)!! / |x|^{2N+3})
  ```

  where `(2n-1)!! = (2n-1) (2n-3) (2n-5) ...  (3) (1)` is a
  [double-factorial](https://en.wikipedia.org/wiki/Double_factorial).

  Args:
    x: `Tensor` of type `float32`, `float64`.
    series_order: Positive Python `integer`. Maximum depth to
      evaluate the asymptotic expansion. This is the `N` above.
    name: Python string. A name for the operation (default="log_ndtr").

  Returns:
    log_ndtr: `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x.dtype` is not handled.
    TypeError: if `series_order` is a not Python `integer.`
    ValueError:  if `series_order` is not in `[0, 30]`.
  """
  if not isinstance(series_order, int):
    raise TypeError("series_order must be a Python integer.")
  if series_order < 0:
    raise ValueError("series_order must be non-negative.")
  if series_order > 30:
    raise ValueError("series_order must be <= 30.")

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")

    if dtype_util.base_equal(x.dtype, tf.float64):
      lower_segment = LOGNDTR_FLOAT64_LOWER
      upper_segment = LOGNDTR_FLOAT64_UPPER
    elif dtype_util.base_equal(x.dtype, tf.float32):
      lower_segment = LOGNDTR_FLOAT32_LOWER
      upper_segment = LOGNDTR_FLOAT32_UPPER
    else:
      raise TypeError("x.dtype=%s is not supported." % x.dtype)

    # The basic idea here was ported from:
    #   https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
    # We copy the main idea, with a few changes
    # * For x >> 1, and X ~ Normal(0, 1),
    #     Log[P[X < x]] = Log[1 - P[X < -x]] approx -P[X < -x],
    #     which extends the range of validity of this function.
    # * We use one fixed series_order for all of 'x', rather than adaptive.
    # * Our docstring properly reflects that this is an asymptotic series, not a
    #   Taylor series. We also provided a correct bound on the remainder.
    # * We need to use the max/min in the _log_ndtr_lower arg to avoid nan when
    #   x=0. This happens even though the branch is unchosen because when x=0
    #   the gradient of a select involves the calculation 1*dy+0*(-inf)=nan
    #   regardless of whether dy is finite. Note that the minimum is a NOP if
    #   the branch is chosen.
    return tf.where(
        tf.greater(x, upper_segment),
        -_ndtr(-x),  # log(1-x) ~= -x, x << 1
        tf.where(
            tf.greater(x, lower_segment),
            tf.math.log(_ndtr(tf.maximum(x, lower_segment))),
            _log_ndtr_lower(tf.minimum(x, lower_segment), series_order)))


Instruction:
Make a summary of the following Python 3 code
def _log_ndtr_lower(x, series_order):
  x_2 = tf.square(x)
  # Log of the term multiplying (1 + sum)
  log_scale = -0.5 * x_2 - tf.math.log(-x) - 0.5 * np.log(2. * np.pi)
  return log_scale + tf.math.log(_log_ndtr_asymptotic_series(x, series_order))

Response: Asymptotic expansion version of Log [ cdf x ] appropriate for x << - 1.


Instruction:
Can you write a function in Python 3 where it
calculates the asymptotic series used in log_ndtr.

Response:
def _log_ndtr_asymptotic_series(x, series_order):
  """Calculates the asymptotic series used in log_ndtr."""
  npdt = dtype_util.as_numpy_dtype(x.dtype)
  if series_order <= 0:
    return npdt(1)
  x_2 = tf.square(x)
  even_sum = tf.zeros_like(x)
  odd_sum = tf.zeros_like(x)
  x_2n = x_2  # Start with x^{2*1} = x^{2*n} with n = 1.
  for n in range(1, series_order + 1):
    y = npdt(_double_factorial(2 * n - 1)) / x_2n
    if n % 2:
      odd_sum += y
    else:
      even_sum += y
    x_2n *= x_2
  return 1. + even_sum - odd_sum


Instruction:
How would you implement a function in Python 3 that
inverses function for erf error function.

Response:
def erfinv(x, name="erfinv"):
  """The inverse function for erf, the error function.

  Args:
    x: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="erfinv").

  Returns:
    x: `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x` is not floating-type.
  """

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")
    if dtype_util.as_numpy_dtype(x.dtype) not in [np.float32, np.float64]:
      raise TypeError("x.dtype={} is not handled, see docstring for supported "
                      "types.".format(dtype_util.name(x.dtype)))
    return ndtri((x + 1.) / 2.) / np.sqrt(2.)


Instruction:
Make a summary of the following Python 3 code
def log_cdf_laplace(x, name="log_cdf_laplace"):

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")

    # For x < 0, L(x) = 0.5 * exp{x} exactly, so Log[L(x)] = log(0.5) + x.
    lower_solution = -np.log(2.) + x

    # safe_exp_neg_x = exp{-x} for x > 0, but is
    # bounded above by 1, which avoids
    #   log[1 - 1] = -inf for x = log(1/2), AND
    #   exp{-x} --> inf, for x << -1
    safe_exp_neg_x = tf.exp(-tf.abs(x))

    # log1p(z) = log(1 + z) approx z for |z| << 1. This approxmation is used
    # internally by log1p, rather than being done explicitly here.
    upper_solution = tf.math.log1p(-0.5 * safe_exp_neg_x)

    return tf.where(x < 0., lower_solution, upper_solution)

Response: Log Laplace distribution function.


Instruction:
How would you explain what the following Python 3 function does
def text_messages_joint_log_prob(count_data, lambda_1, lambda_2, tau):
  alpha = (1. / tf.reduce_mean(input_tensor=count_data))
  rv_lambda = tfd.Exponential(rate=alpha)

  rv_tau = tfd.Uniform()

  lambda_ = tf.gather(
      [lambda_1, lambda_2],
      indices=tf.cast(
          tau * tf.cast(tf.size(input=count_data), dtype=tf.float32) <= tf.cast(
              tf.range(tf.size(input=count_data)), dtype=tf.float32),
          dtype=tf.int32))
  rv_observation = tfd.Poisson(rate=lambda_)

  return (rv_lambda.log_prob(lambda_1) + rv_lambda.log_prob(lambda_2) +
          rv_tau.log_prob(tau) +
          tf.reduce_sum(input_tensor=rv_observation.log_prob(count_data)))

Response: Joint log probability function.


Instruction:
Can you create a Python 3 function that
runs HMC on the text - messages unnormalized posterior.

Response:
def benchmark_text_messages_hmc(
    num_results=int(3e3),
    num_burnin_steps=int(3e3),
    num_leapfrog_steps=3):
  """Runs HMC on the text-messages unnormalized posterior."""

  if not tf.executing_eagerly():
    tf.compat.v1.reset_default_graph()

  # Build a static, pretend dataset.
  count_data = tf.cast(
      tf.concat(
          [tfd.Poisson(rate=15.).sample(43),
           tfd.Poisson(rate=25.).sample(31)],
          axis=0),
      dtype=tf.float32)
  if tf.executing_eagerly():
    count_data = count_data.numpy()
  else:
    with tf.compat.v1.Session():
      count_data = count_data.eval()

  # Define a closure over our joint_log_prob.
  def unnormalized_log_posterior(lambda1, lambda2, tau):
    return text_messages_joint_log_prob(count_data, lambda1, lambda2, tau)

  if tf.executing_eagerly():
    sample_chain = tf.function(tfp.mcmc.sample_chain)
  else:
    sample_chain = tfp.mcmc.sample_chain

  # Initialize the step_size. (It will be automatically adapted.)
  step_size = tf.compat.v2.Variable(
      name='step_size',
      initial_value=tf.constant(0.05, dtype=tf.float32),
      trainable=False)

  def computation():
    """The benchmark computation."""

    initial_chain_state = [
        tf.constant(count_data.mean(), name='init_lambda1'),
        tf.constant(count_data.mean(), name='init_lambda2'),
        tf.constant(0.5, name='init_tau'),
    ]

    unconstraining_bijectors = [
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Sigmoid(),   # Maps [0,1] to R.
    ]

    _, kernel_results = sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=initial_chain_state,
        kernel=tfp.mcmc.TransformedTransitionKernel(
            inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
                target_log_prob_fn=unnormalized_log_posterior,
                num_leapfrog_steps=num_leapfrog_steps,
                step_size=step_size,
                step_size_update_fn=
                tfp.mcmc.make_simple_step_size_update_policy(num_burnin_steps),
                state_gradients_are_stopped=True),
            bijector=unconstraining_bijectors))

    return kernel_results.inner_results.is_accepted

  # Let's force evaluation of graph to ensure build time is not part of our time
  # trial.
  is_accepted_tensor = computation()
  if not tf.executing_eagerly():
    session = tf.compat.v1.Session()
    session.run(tf.compat.v1.global_variables_initializer())
    session.run(is_accepted_tensor)

  start_time = time.time()
  if tf.executing_eagerly():
    is_accepted = computation()
  else:
    is_accepted = session.run(is_accepted_tensor)
  wall_time = time.time() - start_time

  num_accepted = np.sum(is_accepted)
  acceptance_rate = np.float32(num_accepted) / np.float32(num_results)

  return dict(
      iters=(num_results + num_burnin_steps) * num_leapfrog_steps,
      extras={'acceptance_rate': acceptance_rate},
      wall_time=wall_time)


Instruction:
How would you code a function in Python 3 to
return True if the given index_points would yield a univariate marginal.

Response:
def _is_univariate_marginal(self, index_points):
    """True if the given index_points would yield a univariate marginal.

    Args:
      index_points: the set of index set locations at which to compute the
      marginal Gaussian distribution. If this set is of size 1, the marginal is
      univariate.

    Returns:
      is_univariate: Boolean indicating whether the marginal is univariate or
      multivariate. In the case of dynamic shape in the number of index points,
      defaults to "multivariate" since that's the best we can do.
    """
    num_index_points = tf.compat.dimension_value(
        index_points.shape[-(self.kernel.feature_ndims + 1)])
    if num_index_points is None:
      warnings.warn(
          'Unable to detect statically whether the number of index_points is '
          '1. As a result, defaulting to treating the marginal GP at '
          '`index_points` as a multivariate Gaussian. This makes some methods, '
          'like `cdf` unavailable.')
    return num_index_points == 1


Instruction:
How would you code a function in Python 3 to
compute the marginal distribution over function values at index_points.

Response:
def get_marginal_distribution(self, index_points=None):
    """Compute the marginal of this GP over function values at `index_points`.

    Args:
      index_points: `float` `Tensor` representing finite (batch of) vector(s) of
        points in the index set over which the GP is defined. Shape has the form
        `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e` is the number
        (size) of index points in each batch. Ultimately this distribution
        corresponds to a `e`-dimensional multivariate normal. The batch shape
        must be broadcastable with `kernel.batch_shape` and any batch dims
        yielded by `mean_fn`.

    Returns:
      marginal: a `Normal` or `MultivariateNormalLinearOperator` distribution,
        according to whether `index_points` consists of one or many index
        points, respectively.
    """
    with self._name_scope('get_marginal_distribution'):
      # TODO(cgs): consider caching the result here, keyed on `index_points`.
      index_points = self._get_index_points(index_points)
      covariance = self._compute_covariance(index_points)
      loc = self._mean_fn(index_points)
      # If we're sure the number of index points is 1, we can just construct a
      # scalar Normal. This has computational benefits and supports things like
      # CDF that aren't otherwise straightforward to provide.
      if self._is_univariate_marginal(index_points):
        scale = tf.sqrt(covariance)
        # `loc` has a trailing 1 in the shape; squeeze it.
        loc = tf.squeeze(loc, axis=-1)
        return normal.Normal(
            loc=loc,
            scale=scale,
            validate_args=self._validate_args,
            allow_nan_stats=self._allow_nan_stats,
            name='marginal_distribution')
      else:
        scale = tf.linalg.LinearOperatorLowerTriangular(
            tf.linalg.cholesky(_add_diagonal_shift(covariance, self.jitter)),
            is_non_singular=True,
            name='GaussianProcessScaleLinearOperator')
        return mvn_linear_operator.MultivariateNormalLinearOperator(
            loc=loc,
            scale=scale,
            validate_args=self._validate_args,
            allow_nan_stats=self._allow_nan_stats,
            name='marginal_distribution')


Instruction:
Can you write a function in Python 3 where it
returns the index_points argument if given otherwise the class member .

Response:
def _get_index_points(self, index_points=None):
    """Return `index_points` if not None, else `self._index_points`.

    Args:
      index_points: if given, this is what is returned; else,
      `self._index_points`

    Returns:
      index_points: the given arg, if not None, else the class member
      `self._index_points`.

    Rases:
      ValueError: if `index_points` and `self._index_points` are both `None`.
    """
    if self._index_points is None and index_points is None:
      raise ValueError(
          'This GaussianProcess instance was not instantiated with a value for '
          'index_points. One must therefore be provided when calling sample, '
          'log_prob, and other such methods. In particular, one can\'t compute '
          'KL divergences to/from an instance of `GaussianProccess` with '
          'unspecified `index_points` directly. Instead, use the '
          '`get_marginal_distribution` function, which takes `index_points` as '
          'an argument and returns a `Normal` or '
          '`MultivariateNormalLinearOperator` instance, whose KL can be '
          'computed.')
    return index_points if index_points is not None else self._index_points


Instruction:
Given the following Python 3 function, write the documentation
def _logsum_expbig_minus_expsmall(big, small):
  with tf.name_scope("logsum_expbig_minus_expsmall"):
    return tf.math.log1p(-tf.exp(small - big)) + big

Response: Stable evaluation of log [ expbig - expsmall ).


Instruction:
Implement a Python 3 function for
computing log_prob using log survival_function and cdf together.

Response:
def _log_prob_with_logsf_and_logcdf(self, y):
    """Compute log_prob(y) using log survival_function and cdf together."""
    # There are two options that would be equal if we had infinite precision:
    # Log[ sf(y - 1) - sf(y) ]
    #   = Log[ exp{logsf(y - 1)} - exp{logsf(y)} ]
    # Log[ cdf(y) - cdf(y - 1) ]
    #   = Log[ exp{logcdf(y)} - exp{logcdf(y - 1)} ]
    logsf_y = self.log_survival_function(y)
    logsf_y_minus_1 = self.log_survival_function(y - 1)
    logcdf_y = self.log_cdf(y)
    logcdf_y_minus_1 = self.log_cdf(y - 1)

    # Important:  Here we use select in a way such that no input is inf, this
    # prevents the troublesome case where the output of select can be finite,
    # but the output of grad(select) will be NaN.

    # In either case, we are doing Log[ exp{big} - exp{small} ]
    # We want to use the sf items precisely when we are on the right side of the
    # median, which occurs when logsf_y < logcdf_y.
    big = tf.where(logsf_y < logcdf_y, logsf_y_minus_1, logcdf_y)
    small = tf.where(logsf_y < logcdf_y, logsf_y, logcdf_y_minus_1)

    return _logsum_expbig_minus_expsmall(big, small)


Instruction:
How would you code a function in Python 3 to
create an IAF bijector.

Response:
def make_iaf_stack(total_event_size,
                   num_hidden_layers=2,
                   seed=None,
                   dtype=tf.float32):
  """Creates an stacked IAF bijector.

  This bijector operates on vector-valued events.

  Args:
    total_event_size: Number of dimensions to operate over.
    num_hidden_layers: How many hidden layers to use in each IAF.
    seed: Random seed for the initializers.
    dtype: DType for the variables.

  Returns:
    bijector: The created bijector.
  """

  seed = tfd.SeedStream(seed, 'make_iaf_stack')

  def make_iaf():
    """Create an IAF."""
    initializer = tf.compat.v2.keras.initializers.VarianceScaling(
        2 * 0.01, seed=seed() % (2**31 - 1))

    made = tfb.AutoregressiveLayer(
        params=2,
        event_shape=[total_event_size],
        hidden_units=[total_event_size] * num_hidden_layers,
        activation=tf.nn.elu,
        kernel_initializer=initializer,
        dtype=dtype)

    def shift_and_scale(x):
      # TODO(siege): Something is losing the static shape.
      x.set_shape(
          x.shape.merge_with([None] * (x.shape.ndims - 1) + [total_event_size]))
      return tf.unstack(made(x), num=2, axis=-1)

    return tfb.Invert(tfb.MaskedAutoregressiveFlow(shift_and_scale))

  def make_swap():
    """Create an swap."""
    permutation = list(reversed(range(total_event_size)))
    return tfb.Permute(permutation)

  bijector = make_iaf()
  bijector = make_swap()(bijector)
  bijector = make_iaf()(bijector)
  bijector = make_swap()(bijector)
  bijector = make_iaf()(bijector)
  bijector = make_swap()(bijector)

  return bijector


Instruction:
Write a Python 3 script to
run one iteration of NeuTra.

Response:
def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of NeuTra.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
    """

    @tfp.mcmc.internal.util.make_innermost_setter
    def set_num_leapfrog_steps(kernel_results, num_leapfrog_steps):
      return kernel_results._replace(
          accepted_results=kernel_results.accepted_results._replace(
              num_leapfrog_steps=num_leapfrog_steps))

    step_size = previous_kernel_results.new_step_size
    previous_kernel_results = set_num_leapfrog_steps(
        previous_kernel_results, self._num_leapfrog_steps(step_size))

    new_state, kernel_results = self._kernel.one_step(
        self._flatten_state(current_state), previous_kernel_results)
    return self._unflatten_state(new_state), kernel_results


Instruction:
Write a Python 3 function that can
train the bijector and creates initial kernel_results.

Response:
def bootstrap_results(self, state):
    """Trains the bijector and creates initial `previous_kernel_results`.

    The supplied `state` is only used to determine the number of chains to run
    in parallel_iterations

    Args:
      state: `Tensor` or Python `list` of `Tensor`s representing the initial
        state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*state))`.

    Returns:
      kernel_results: Instance of
        `UncalibratedHamiltonianMonteCarloKernelResults` inside
        `MetropolisHastingsResults` inside `TransformedTransitionKernelResults`
        inside `SimpleStepSizeAdaptationResults`.
    """

    def loss():
      q = self._flattened_variational_distribution()
      # TODO(siege): How to seed this?
      samples = q.sample(self.train_batch_size)
      return tf.reduce_mean(
          input_tensor=q.log_prob(samples) -
          self._flattened_target_log_prob(samples),
          axis=-1)

    lr = tf.convert_to_tensor(value=self.learning_rate, dtype=self._dtype)
    dtype = lr.dtype

    learning_rate = tf.compat.v2.optimizers.schedules.PiecewiseConstantDecay(
        list(self.num_train_steps *
             np.array([0.2, 0.8]).astype(dtype.as_numpy_dtype())),
        [lr, lr * 0.1, lr * 0.01])

    opt = tf.compat.v2.optimizers.Adam(learning_rate)

    @tf.function(autograph=False)
    def train_step():
      with tf.GradientTape() as tape:
        loss_val = loss()
      vals = tape.watched_variables()
      grads = tape.gradient(loss_val, vals)
      grads_and_vals = list(zip(grads, vals))
      opt.apply_gradients(grads_and_vals)
      return loss_val

    for step in range(self.num_train_steps):
      loss_val = train_step()
      tf.debugging.assert_all_finite(
          loss_val, 'NeuTra loss is NaN at step {}'.format(step))
      if self.train_debug_fn:
        # pylint: disable=not-callable
        self.train_debug_fn(self, step, loss_val)

    state_parts = tf.nest.flatten(state)
    flat_state_shapes = tf.nest.flatten(self.state_shape)
    batch_shape = tf.shape(input=state_parts[0])[:-flat_state_shapes[0].ndims]

    return self._kernel.bootstrap_results(
        self._flattened_variational_distribution().sample(
            batch_shape, seed=self.seed))


Instruction:
Explain what the following Python 3 code does
def mean_log_det(self, name="mean_log_det"):
    with self._name_scope(name):
      return (self._multi_digamma(0.5 * self.df, self.dimension) +
              self.dimension * math.log(2.) +
              2 * self.scale_operator.log_abs_determinant())

Response: Computes the mean log det of X under this Wishart distribution.


Instruction:
Write a Python 3 function for
computing the log normalizing constant log ( Z.

Response:
def log_normalization(self, name="log_normalization"):
    """Computes the log normalizing constant, log(Z)."""
    with self._name_scope(name):
      return (self.df * self.scale_operator.log_abs_determinant() +
              0.5 * self.df * self.dimension * math.log(2.) +
              self._multi_lgamma(0.5 * self.df, self.dimension))


Instruction:
Can you write a function in Python 3 where it
creates a sequence used in multivariate gamma ; shape = shape a + p.

Response:
def _multi_gamma_sequence(self, a, p, name="multi_gamma_sequence"):
    """Creates sequence used in multivariate (di)gamma; shape = shape(a)+[p]."""
    with self._name_scope(name):
      # Linspace only takes scalars, so we'll add in the offset afterwards.
      seq = tf.linspace(
          tf.constant(0., dtype=self.dtype), 0.5 - 0.5 * p, tf.cast(
              p, tf.int32))
      return seq + tf.expand_dims(a, [-1])


Instruction:
Here you have a function in Python 3, explain what it does
def _multi_lgamma(self, a, p, name="multi_lgamma"):
    with self._name_scope(name):
      seq = self._multi_gamma_sequence(a, p)
      return (0.25 * p * (p - 1.) * math.log(math.pi) +
              tf.reduce_sum(input_tensor=tf.math.lgamma(seq), axis=[-1]))

Response: Computes the log multivariate gamma function ; log ( Gamma_p a ).


Instruction:
Given the following Python 3 function, write the documentation
def _multi_digamma(self, a, p, name="multi_digamma"):
    with self._name_scope(name):
      seq = self._multi_gamma_sequence(a, p)
      return tf.reduce_sum(input_tensor=tf.math.digamma(seq), axis=[-1])

Response: Computes the multivariate digamma function ; Psi_p a.


Instruction:
Here you have a function in Python 3, explain what it does
def _outer_squared_difference(x, y):
  z = x - y
  return z[..., tf.newaxis, :] * z[..., tf.newaxis]

Response: Convenience function analogous to tf. squared_difference.


Instruction:
Can you tell what is the following Python 3 function doing
def _value_and_batch_jacobian(f, x):
  if tf.executing_eagerly():
    with tf.GradientTape() as tape:
      tape.watch(x)
      value = f(x)
    batch_jacobian = tape.batch_jacobian(value, x)
  else:
    value = f(x)
    batch_jacobian = gradients.batch_jacobian(value, x)
  return value, batch_jacobian

Response: Enables uniform interface to value and batch jacobian calculation.


Instruction:
Write a Python 3 script for
disabling computation of the second derivatives for a tensor.

Response:
def _prevent_2nd_derivative(x):
  """Disables computation of the second derivatives for a tensor.

  NB: you need to apply a non-identity function to the output tensor for the
  exception to be raised.

  Arguments:
    x: A tensor.

  Returns:
    A tensor with the same value and the same derivative as x, but that raises
    LookupError when trying to compute the second derivatives.
  """
  def grad(dy):
    return array_ops.prevent_gradient(
        dy, message="Second derivative is not implemented.")

  return tf.identity(x), grad


Instruction:
Implement a Python 3 function for
adding reparameterization (pathwise) gradients to samples of the mixture. Implicit reparameterization gradients are dx/dphi = -(d transform(x, phi) / dx)^-1 * d transform(x, phi) / dphi, where transform(x, phi) is distributional transform that removes all parameters from samples x. We implement them by replacing x with -stop_gradient(d transform(x, phi) / dx)^-1 * transform(x, phi)] for the backward pass (gradient computation). The derivative of this quantity w.r.t. phi is then the implicit reparameterization gradient. Note that this replaces the gradients w.r.t. both the mixture distribution parameters and components distributions parameters. Limitations: 1. Fundamental: components must be fully reparameterized. 2. Distributional transform is currently only implemented for factorized components. 3. Distributional transform currently only works for known rank of the batch tensor. Arguments: x: Sample of mixture distribution Returns: Tensor with same value as x, but with reparameterization gradients

Response:
def _reparameterize_sample(self, x):
    """Adds reparameterization (pathwise) gradients to samples of the mixture.

    Implicit reparameterization gradients are
       dx/dphi = -(d transform(x, phi) / dx)^-1 * d transform(x, phi) / dphi,
    where transform(x, phi) is distributional transform that removes all
    parameters from samples x.

    We implement them by replacing x with
      -stop_gradient(d transform(x, phi) / dx)^-1 * transform(x, phi)]
    for the backward pass (gradient computation).
    The derivative of this quantity w.r.t. phi is then the implicit
    reparameterization gradient.
    Note that this replaces the gradients w.r.t. both the mixture
    distribution parameters and components distributions parameters.

    Limitations:
      1. Fundamental: components must be fully reparameterized.
      2. Distributional transform is currently only implemented for
        factorized components.
      3. Distributional transform currently only works for known rank of the
        batch tensor.

    Arguments:
      x: Sample of mixture distribution

    Returns:
      Tensor with same value as x, but with reparameterization gradients
    """
    # Remove the existing gradients of x wrt parameters of the components.
    x = tf.stop_gradient(x)

    x_2d_shape = [-1, self._event_size]  # [S*prod(B), prod(E)]

    # Perform distributional transform of x in [S, B, E] shape,
    # but have Jacobian of size [S*prod(B), prod(E), prod(E)].
    def reshaped_distributional_transform(x_2d):
      return tf.reshape(
          self._distributional_transform(tf.reshape(x_2d, tf.shape(input=x))),
          x_2d_shape)

    # transform_2d: [S*prod(B), prod(E)]
    # jacobian: [S*prod(B), prod(E), prod(E)]
    transform_2d, jacobian = _value_and_batch_jacobian(
        reshaped_distributional_transform, tf.reshape(x, x_2d_shape))

    # We only provide the first derivative; the second derivative computed by
    # autodiff would be incorrect, so we raise an error if it is requested.
    transform_2d = _prevent_2nd_derivative(transform_2d)

    # Compute [- stop_gradient(jacobian)^-1 * transform] by solving a linear
    # system. The Jacobian is lower triangular because the distributional
    # transform for i-th event dimension does not depend on the next
    # dimensions.
    surrogate_x_2d = -tf.linalg.triangular_solve(
        tf.stop_gradient(jacobian), tf.expand_dims(transform_2d, axis=-1),
        lower=True)  # [S*prod(B), prod(E), 1]
    surrogate_x = tf.reshape(surrogate_x_2d, tf.shape(input=x))

    # Replace gradients of x with gradients of surrogate_x, but keep the value.
    return x + (surrogate_x - tf.stop_gradient(surrogate_x))


Instruction:
Here you have a function in Python 3, explain what it does
def _distributional_transform(self, x):

    if tensorshape_util.rank(x.shape) is None:
      # tf.nn.softmax raises an error when applied to inputs of undefined rank.
      raise ValueError("Distributional transform does not support inputs of "
                       "undefined rank.")

    # Obtain factorized components distribution and assert that it's
    # a scalar distribution.
    if isinstance(self._components_distribution, independent.Independent):
      univariate_components = self._components_distribution.distribution
    else:
      univariate_components = self._components_distribution

    with tf.control_dependencies([
        assert_util.assert_equal(
            univariate_components.is_scalar_event(),
            True,
            message="`univariate_components` must have scalar event")
    ]):
      x_padded = self._pad_sample_dims(x)  # [S, B, 1, E]
      log_prob_x = univariate_components.log_prob(x_padded)  # [S, B, k, E]
      cdf_x = univariate_components.cdf(x_padded)  # [S, B, k, E]

      # log prob_k (x_1, ..., x_i-1)
      cumsum_log_prob_x = tf.reshape(
          tf.math.cumsum(
              # [S*prod(B)*k, prod(E)]
              tf.reshape(log_prob_x, [-1, self._event_size]),
              exclusive=True,
              axis=-1),
          tf.shape(input=log_prob_x))  # [S, B, k, E]

      logits_mix_prob = distribution_utils.pad_mixture_dimensions(
          self.mixture_distribution.logits, self, self.mixture_distribution,
          self._event_ndims)  # [B, k, 1]

      # Logits of the posterior weights: log w_k + log prob_k (x_1, ..., x_i-1)
      log_posterior_weights_x = logits_mix_prob + cumsum_log_prob_x

      component_axis = tensorshape_util.rank(x.shape) - self._event_ndims
      posterior_weights_x = tf.nn.softmax(log_posterior_weights_x,
                                          axis=component_axis)
      return tf.reduce_sum(
          input_tensor=posterior_weights_x * cdf_x, axis=component_axis)

Response: Performs the distributional transform of the mixture samples.


Instruction:
Here you have a function in Python 3, explain what it does
def _split_covariance_into_marginals(covariance, block_sizes):
  start_dim = 0
  marginals = []
  for size in block_sizes:
    end_dim = start_dim + size
    marginals.append(covariance[..., start_dim:end_dim, start_dim:end_dim])
    start_dim = end_dim
  return marginals

Response: Split a covariance matrix into block - diagonal marginals of given sizes.


Instruction:
Write a Python 3 function that can
decompose an observed time series into contributions from each component. This method decomposes a time series according to the posterior represention of a structural time series model. In particular, it: - Computes the posterior marginal mean and covariances over the additive model's latent space. - Decomposes the latent posterior into the marginal blocks for each model component. - Maps the per-component latent posteriors back through each component's observation model, to generate the time series modeled by that component. Args: model: An instance of `tfp.sts.Sum` representing a structural time series model. observed_time_series: `float` `Tensor` of shape `batch_shape + [num_timesteps, 1]` (omitting the trailing unit dimension is also supported when `num_timesteps > 1`), specifying an observed time series. May optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes a mask `Tensor` to specify timesteps with missing observations. parameter_samples: Python `list` of `Tensors` representing posterior samples of model parameters, with shapes `[concat([ [num_posterior_draws], param.prior.batch_shape, param.prior.event_shape]) for param in model.parameters]`. This may optionally also be a map (Python `dict`) of parameter names to `Tensor` values. Returns: component_dists: A `collections.OrderedDict` instance mapping component StructuralTimeSeries instances (elements of `model.components`) to `tfd.Distribution` instances representing the posterior marginal distributions on the process modeled by each component. Each distribution has batch shape matching that of `posterior_means`/`posterior_covs`, and event shape of `[num_timesteps]`. #### Examples Suppose we've built a model and fit it to data: ```python day_of_week = tfp.sts.Seasonal( num_seasons=7, observed_time_series=observed_time_series, name='day_of_week') local_linear_trend = tfp.sts.LocalLinearTrend( observed_time_series=observed_time_series, name='local_linear_trend') model = tfp.sts.Sum(components=[day_of_week, local_linear_trend], observed_time_series=observed_time_series) num_steps_forecast = 50 samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series) ``` To extract the contributions of individual components, pass the time series and sampled parameters into `decompose_by_component`: ```python component_dists = decompose_by_component( model, observed_time_series=observed_time_series, parameter_samples=samples) # Component mean and stddev have shape `[len(observed_time_series)]`. day_of_week_effect_mean = component_dists[day_of_week].mean() day_of_week_effect_stddev = component_dists[day_of_week].stddev() ``` Using the component distributions, we can visualize the uncertainty for each component: ``` from matplotlib import pylab as plt num_components = len(component_dists) xs = np.arange(len(observed_time_series)) fig = plt.figure(figsize=(12, 3 * num_components)) for i, (component, component_dist) in enumerate(component_dists.items()): # If in graph mode, replace `.numpy()` with `.eval()` or `sess.run()`. component_mean = component_dist.mean().numpy() component_stddev = component_dist.stddev().numpy() ax = fig.add_subplot(num_components, 1, 1 + i) ax.plot(xs, component_mean, lw=2) ax.fill_between(xs, component_mean - 2 * component_stddev, component_mean + 2 * component_stddev, alpha=0.5) ax.set_title(component.name) ```

Response:
def decompose_by_component(model, observed_time_series, parameter_samples):
  """Decompose an observed time series into contributions from each component.

  This method decomposes a time series according to the posterior represention
  of a structural time series model. In particular, it:
    - Computes the posterior marginal mean and covariances over the additive
      model's latent space.
    - Decomposes the latent posterior into the marginal blocks for each
      model component.
    - Maps the per-component latent posteriors back through each component's
      observation model, to generate the time series modeled by that component.

  Args:
    model: An instance of `tfp.sts.Sum` representing a structural time series
      model.
    observed_time_series: `float` `Tensor` of shape
      `batch_shape + [num_timesteps, 1]` (omitting the trailing unit dimension
      is also supported when `num_timesteps > 1`), specifying an observed time
      series. May optionally be an instance of `tfp.sts.MaskedTimeSeries`, which
      includes a mask `Tensor` to specify timesteps with missing observations.
    parameter_samples: Python `list` of `Tensors` representing posterior
      samples of model parameters, with shapes `[concat([
      [num_posterior_draws], param.prior.batch_shape,
      param.prior.event_shape]) for param in model.parameters]`. This may
      optionally also be a map (Python `dict`) of parameter names to
      `Tensor` values.
  Returns:
    component_dists: A `collections.OrderedDict` instance mapping
      component StructuralTimeSeries instances (elements of `model.components`)
      to `tfd.Distribution` instances representing the posterior marginal
      distributions on the process modeled by each component. Each distribution
      has batch shape matching that of `posterior_means`/`posterior_covs`, and
      event shape of `[num_timesteps]`.

  #### Examples

  Suppose we've built a model and fit it to data:

  ```python
    day_of_week = tfp.sts.Seasonal(
        num_seasons=7,
        observed_time_series=observed_time_series,
        name='day_of_week')
    local_linear_trend = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series,
        name='local_linear_trend')
    model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                        observed_time_series=observed_time_series)

    num_steps_forecast = 50
    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
  ```

  To extract the contributions of individual components, pass the time series
  and sampled parameters into `decompose_by_component`:

  ```python
    component_dists = decompose_by_component(
      model,
      observed_time_series=observed_time_series,
      parameter_samples=samples)

    # Component mean and stddev have shape `[len(observed_time_series)]`.
    day_of_week_effect_mean = component_dists[day_of_week].mean()
    day_of_week_effect_stddev = component_dists[day_of_week].stddev()
  ```

  Using the component distributions, we can visualize the uncertainty for
  each component:

  ```
  from matplotlib import pylab as plt
  num_components = len(component_dists)
  xs = np.arange(len(observed_time_series))
  fig = plt.figure(figsize=(12, 3 * num_components))
  for i, (component, component_dist) in enumerate(component_dists.items()):

    # If in graph mode, replace `.numpy()` with `.eval()` or `sess.run()`.
    component_mean = component_dist.mean().numpy()
    component_stddev = component_dist.stddev().numpy()

    ax = fig.add_subplot(num_components, 1, 1 + i)
    ax.plot(xs, component_mean, lw=2)
    ax.fill_between(xs,
                    component_mean - 2 * component_stddev,
                    component_mean + 2 * component_stddev,
                    alpha=0.5)
    ax.set_title(component.name)
  ```

  """

  with tf.compat.v1.name_scope('decompose_by_component',
                               values=[observed_time_series]):
    [
        observed_time_series,
        is_missing
    ] = sts_util.canonicalize_observed_time_series_with_mask(
        observed_time_series)

    # Run smoothing over the training timesteps to extract the
    # posterior on latents.
    num_timesteps = dist_util.prefer_static_value(
        tf.shape(input=observed_time_series))[-2]
    ssm = model.make_state_space_model(num_timesteps=num_timesteps,
                                       param_vals=parameter_samples)
    posterior_means, posterior_covs = ssm.posterior_marginals(
        observed_time_series, mask=is_missing)

    return _decompose_from_posterior_marginals(
        model, posterior_means, posterior_covs, parameter_samples)


Instruction:
How would you explain what the following Python 3 function does
def decompose_forecast_by_component(model, forecast_dist, parameter_samples):

  with tf.compat.v1.name_scope('decompose_forecast_by_component'):
    try:
      forecast_lgssm = forecast_dist.components_distribution
      forecast_latent_mean, _ = forecast_lgssm._joint_mean()  # pylint: disable=protected-access
      forecast_latent_covs, _ = forecast_lgssm._joint_covariances()  # pylint: disable=protected-access
    except AttributeError as e:
      raise ValueError(
          'Forecast distribution must be a MixtureSameFamily of'
          'LinearGaussianStateSpaceModel distributions, such as returned by'
          '`tfp.sts.forecast()`. (saw exception: {})'.format(e))

    # Since `parameter_samples` will have sample shape `[num_posterior_draws]`,
    # we need to move the `num_posterior_draws` dimension of the forecast
    # moments from the trailing batch dimension, where it's currently put by
    # `sts.forecast`, back to the leading (sample shape) dimension.
    forecast_latent_mean = dist_util.move_dimension(
        forecast_latent_mean, source_idx=-3, dest_idx=0)
    forecast_latent_covs = dist_util.move_dimension(
        forecast_latent_covs, source_idx=-4, dest_idx=0)

    return _decompose_from_posterior_marginals(
        model, forecast_latent_mean, forecast_latent_covs, parameter_samples)

Response: Decompose a forecast distribution into contributions from each component.


Instruction:
Write a Python 3 script to
convert dense tensor x to SparseTensor.

Response:
def dense_to_sparse(x, ignore_value=None, name=None):
  """Converts dense `Tensor` to `SparseTensor`, dropping `ignore_value` cells.

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

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

  Raises:
    ValueError: when `x`'s rank is `None`.
  """
  # Copied (with modifications) from:
  # tensorflow/contrib/layers/python/ops/sparse_ops.py.
  with tf.compat.v1.name_scope(name, 'dense_to_sparse', [x, ignore_value]):
    x = tf.convert_to_tensor(value=x, name='x')
    if ignore_value is None:
      if x.dtype.base_dtype == tf.string:
        # Exception due to TF strings are converted to numpy objects by default.
        ignore_value = ''
      else:
        ignore_value = x.dtype.as_numpy_dtype(0)
      ignore_value = tf.cast(ignore_value, x.dtype, name='ignore_value')
    indices = tf.where(tf.not_equal(x, ignore_value), name='indices')
    return tf.SparseTensor(
        indices=indices,
        values=tf.gather_nd(x, indices, name='values'),
        dense_shape=tf.shape(input=x, out_type=tf.int64, name='dense_shape'))


Instruction:
Can you generate the documentation for the following Python 3 function
def _operator(attr):
  @functools.wraps(attr)
  def func(a, *args):
    return attr(a.value, *args)
  return func

Response: Defers an operator overload to attr.


Instruction:
Given the following Python 3 function, write the documentation
def sample_shape(self):
    if isinstance(self._sample_shape, tf.Tensor):
      return tf.TensorShape(tf.get_static_value(self._sample_shape))
    return tf.TensorShape(self._sample_shape)

Response: Sample shape of random variable as a TensorShape.


Instruction:
How would you implement a function in Python 3 that
returns the sample_shape of the random variable as a 1 - D Tensor.

Response:
def sample_shape_tensor(self, name="sample_shape_tensor"):
    """Sample shape of random variable as a 1-D `Tensor`.

    Args:
      name: name to give to the op

    Returns:
      sample_shape: `Tensor`.
    """
    with tf.compat.v1.name_scope(name):
      if isinstance(self._sample_shape, tf.Tensor):
        return self._sample_shape
      return tf.convert_to_tensor(
          value=self.sample_shape.as_list(), dtype=tf.int32)


Instruction:
Can you implement a function in Python 3 that
gets tensor that the random variable corresponds to.

Response:
def value(self):
    """Get tensor that the random variable corresponds to."""
    if self._value is None:
      try:
        self._value = self.distribution.sample(self.sample_shape_tensor())
      except NotImplementedError:
        raise NotImplementedError(
            "sample is not implemented for {0}. You must either pass in the "
            "value argument or implement sample for {0}."
            .format(self.distribution.__class__.__name__))
    return self._value


Instruction:
Can you tell what is the following Python 3 function doing
def numpy(self):
    if not isinstance(self.value, ops.EagerTensor):
      raise NotImplementedError("value argument must be a EagerTensor.")

    return self.value.numpy()

Response: Value as NumPy array only available for TF Eager.


Instruction:
Write a Python 3 script for
building a scale-and-shift function using a multi-layer neural network. This will be wrapped in a make_template to ensure the variables are only created once. It takes the `d`-dimensional input x[0:d] and returns the `D-d` dimensional outputs `loc` ("mu") and `log_scale` ("alpha"). The default template does not support conditioning and will raise an exception if `condition_kwargs` are passed to it. To use conditioning in real nvp bijector, implement a conditioned shift/scale template that handles the `condition_kwargs`. Arguments: hidden_layers: Python `list`-like of non-negative integer, scalars indicating the number of units in each hidden layer. Default: `[512, 512]. shift_only: Python `bool` indicating if only the `shift` term shall be computed (i.e. NICE bijector). Default: `False`. activation: Activation function (callable). Explicitly setting to `None` implies a linear activation. name: A name for ops managed by this function. Default: "real_nvp_default_template". *args: `tf.layers.dense` arguments. **kwargs: `tf.layers.dense` keyword arguments. Returns: shift: `Float`-like `Tensor` of shift terms ("mu" in [Papamakarios et al. (2016)][1]). log_scale: `Float`-like `Tensor` of log(scale) terms ("alpha" in [Papamakarios et al. (2016)][1]). Raises: NotImplementedError: if rightmost dimension of `inputs` is unknown prior to graph execution, or if `condition_kwargs` is not empty. #### References [1]: George Papamakarios, Theo Pavlakou, and Iain Murray. Masked Autoregressive Flow for Density Estimation. In _Neural Information Processing Systems_, 2017. https://arxiv.org/abs/1705.07057

Response:
def real_nvp_default_template(hidden_layers,
                              shift_only=False,
                              activation=tf.nn.relu,
                              name=None,
                              *args,  # pylint: disable=keyword-arg-before-vararg
                              **kwargs):
  """Build a scale-and-shift function using a multi-layer neural network.

  This will be wrapped in a make_template to ensure the variables are only
  created once. It takes the `d`-dimensional input x[0:d] and returns the `D-d`
  dimensional outputs `loc` ("mu") and `log_scale` ("alpha").

  The default template does not support conditioning and will raise an
  exception if `condition_kwargs` are passed to it. To use conditioning in
  real nvp bijector, implement a conditioned shift/scale template that
  handles the `condition_kwargs`.

  Arguments:
    hidden_layers: Python `list`-like of non-negative integer, scalars
      indicating the number of units in each hidden layer. Default: `[512, 512].
    shift_only: Python `bool` indicating if only the `shift` term shall be
      computed (i.e. NICE bijector). Default: `False`.
    activation: Activation function (callable). Explicitly setting to `None`
      implies a linear activation.
    name: A name for ops managed by this function. Default:
      "real_nvp_default_template".
    *args: `tf.layers.dense` arguments.
    **kwargs: `tf.layers.dense` keyword arguments.

  Returns:
    shift: `Float`-like `Tensor` of shift terms ("mu" in
      [Papamakarios et al.  (2016)][1]).
    log_scale: `Float`-like `Tensor` of log(scale) terms ("alpha" in
      [Papamakarios et al. (2016)][1]).

  Raises:
    NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
      graph execution, or if `condition_kwargs` is not empty.

  #### References

  [1]: George Papamakarios, Theo Pavlakou, and Iain Murray. Masked
       Autoregressive Flow for Density Estimation. In _Neural Information
       Processing Systems_, 2017. https://arxiv.org/abs/1705.07057
  """

  with tf.compat.v2.name_scope(name or "real_nvp_default_template"):

    def _fn(x, output_units, **condition_kwargs):
      """Fully connected MLP parameterized via `real_nvp_template`."""
      if condition_kwargs:
        raise NotImplementedError(
            "Conditioning not implemented in the default template.")

      if tensorshape_util.rank(x.shape) == 1:
        x = x[tf.newaxis, ...]
        reshape_output = lambda x: x[0]
      else:
        reshape_output = lambda x: x
      for units in hidden_layers:
        x = tf.compat.v1.layers.dense(
            inputs=x,
            units=units,
            activation=activation,
            *args,  # pylint: disable=keyword-arg-before-vararg
            **kwargs)
      x = tf.compat.v1.layers.dense(
          inputs=x,
          units=(1 if shift_only else 2) * output_units,
          activation=None,
          *args,  # pylint: disable=keyword-arg-before-vararg
          **kwargs)
      if shift_only:
        return reshape_output(x), None
      shift, log_scale = tf.split(x, 2, axis=-1)
      return reshape_output(shift), reshape_output(log_scale)

    return tf.compat.v1.make_template("real_nvp_default_template", _fn)


Instruction:
Explain what the following Python 3 code does
def _uniform_unit_norm(dimension, shape, dtype, seed):
  # This works because the Gaussian distribution is spherically symmetric.
  # raw shape: shape + [dimension]
  raw = normal.Normal(
      loc=dtype_util.as_numpy_dtype(dtype)(0),
      scale=dtype_util.as_numpy_dtype(dtype)(1)).sample(
          tf.concat([shape, [dimension]], axis=0), seed=seed())
  unit_norm = raw / tf.norm(tensor=raw, ord=2, axis=-1)[..., tf.newaxis]
  return unit_norm

Response: Returns a batch of points chosen uniformly from the unit hypersphere.


Instruction:
Create a Python 3 function to
replicate the input tensor n times along a new major dimension.

Response:
def _replicate(n, tensor):
  """Replicate the input tensor n times along a new (major) dimension."""
  # TODO(axch) Does this already exist somewhere?  Should it get contributed?
  multiples = tf.concat([[n], tf.ones_like(tensor.shape)], axis=0)
  return tf.tile(tf.expand_dims(tensor, axis=0), multiples)


Instruction:
Given the following Python 3 function, write the documentation
def _sample_n(self, num_samples, seed=None, name=None):
    if self.dimension < 0:
      raise ValueError(
          'Cannot sample negative-dimension correlation matrices.')
    # Notation below: B is the batch shape, i.e., tf.shape(concentration)
    seed = seed_stream.SeedStream(seed, 'sample_lkj')
    with tf.name_scope('sample_lkj' or name):
      if not dtype_util.is_floating(self.concentration.dtype):
        raise TypeError(
            'The concentration argument should have floating type, not '
            '{}'.format(dtype_util.name(self.concentration.dtype)))

      concentration = _replicate(num_samples, self.concentration)
      concentration_shape = tf.shape(input=concentration)
      if self.dimension <= 1:
        # For any dimension <= 1, there is only one possible correlation matrix.
        shape = tf.concat([
            concentration_shape, [self.dimension, self.dimension]], axis=0)
        return tf.ones(shape=shape, dtype=self.concentration.dtype)
      beta_conc = concentration + (self.dimension - 2.) / 2.
      beta_dist = beta.Beta(concentration1=beta_conc, concentration0=beta_conc)

      # Note that the sampler below deviates from [1], by doing the sampling in
      # cholesky space. This does not change the fundamental logic of the
      # sampler, but does speed up the sampling.

      # This is the correlation coefficient between the first two dimensions.
      # This is also `r` in reference [1].
      corr12 = 2. * beta_dist.sample(seed=seed()) - 1.

      # Below we construct the Cholesky of the initial 2x2 correlation matrix,
      # which is of the form:
      # [[1, 0], [r, sqrt(1 - r**2)]], where r is the correlation between the
      # first two dimensions.
      # This is the top-left corner of the cholesky of the final sample.
      first_row = tf.concat([
          tf.ones_like(corr12)[..., tf.newaxis],
          tf.zeros_like(corr12)[..., tf.newaxis]], axis=-1)
      second_row = tf.concat([
          corr12[..., tf.newaxis],
          tf.sqrt(1 - corr12**2)[..., tf.newaxis]], axis=-1)

      chol_result = tf.concat([
          first_row[..., tf.newaxis, :],
          second_row[..., tf.newaxis, :]], axis=-2)

      for n in range(2, self.dimension):
        # Loop invariant: on entry, result has shape B + [n, n]
        beta_conc -= 0.5
        # norm is y in reference [1].
        norm = beta.Beta(
            concentration1=n/2.,
            concentration0=beta_conc
        ).sample(seed=seed())
        # distance shape: B + [1] for broadcast
        distance = tf.sqrt(norm)[..., tf.newaxis]
        # direction is u in reference [1].
        # direction shape: B + [n]
        direction = _uniform_unit_norm(
            n, concentration_shape, self.concentration.dtype, seed)
        # raw_correlation is w in reference [1].
        raw_correlation = distance * direction  # shape: B + [n]

        # This is the next row in the cholesky of the result,
        # which differs from the construction in reference [1].
        # In the reference, the new row `z` = chol_result @ raw_correlation^T
        # = C @ raw_correlation^T (where as short hand we use C = chol_result).
        # We prove that the below equation is the right row to add to the
        # cholesky, by showing equality with reference [1].
        # Let S be the sample constructed so far, and let `z` be as in
        # reference [1]. Then at this iteration, the new sample S' will be
        # [[S z^T]
        #  [z 1]]
        # In our case we have the cholesky decomposition factor C, so
        # we want our new row x (same size as z) to satisfy:
        #  [[S z^T]  [[C 0]    [[C^T  x^T]         [[CC^T  Cx^T]
        #   [z 1]] =  [x k]]    [0     k]]  =       [xC^t   xx^T + k**2]]
        # Since C @ raw_correlation^T = z = C @ x^T, and C is invertible,
        # we have that x = raw_correlation. Also 1 = xx^T + k**2, so k
        # = sqrt(1 - xx^T) = sqrt(1 - |raw_correlation|**2) = sqrt(1 -
        # distance**2).
        new_row = tf.concat(
            [raw_correlation, tf.sqrt(1. - norm[..., tf.newaxis])], axis=-1)

        # Finally add this new row, by growing the cholesky of the result.
        chol_result = tf.concat([
            chol_result,
            tf.zeros_like(chol_result[..., 0][..., tf.newaxis])], axis=-1)

        chol_result = tf.concat(
            [chol_result, new_row[..., tf.newaxis, :]], axis=-2)

      if self.input_output_cholesky:
        return chol_result

      result = tf.matmul(chol_result, chol_result, transpose_b=True)
      # The diagonal for a correlation matrix should always be ones. Due to
      # numerical instability the matmul might not achieve that, so manually set
      # these to ones.
      result = tf.linalg.set_diag(
          result,
          tf.ones(shape=tf.shape(input=result)[:-1], dtype=result.dtype))
      # This sampling algorithm can produce near-PSD matrices on which standard
      # algorithms such as `tf.cholesky` or `tf.linalg.self_adjoint_eigvals`
      # fail. Specifically, as documented in b/116828694, around 2% of trials
      # of 900,000 5x5 matrices (distributed according to 9 different
      # concentration parameter values) contained at least one matrix on which
      # the Cholesky decomposition failed.
      return result

Response: Returns a Tensor of samples from an LKJ distribution.


Instruction:
Explain what the following Python 3 code does
def _log_unnorm_prob(self, x, name=None):
    with tf.name_scope(name or 'log_unnorm_prob_lkj'):
      x = tf.convert_to_tensor(value=x, name='x')
      # The density is det(matrix) ** (concentration - 1).
      # Computing the determinant with `logdet` is usually fine, since
      # correlation matrices are Hermitian and PSD. But in some cases, for a
      # PSD matrix whose eigenvalues are close to zero, `logdet` raises an error
      # complaining that it is not PSD. The root cause is the computation of the
      # cholesky decomposition in `logdet`. Hence, we use the less efficient but
      # more robust `slogdet` which does not use `cholesky`.
      #
      # An alternative would have been to check allow_nan_stats and use
      #   eigenvalues = tf.linalg.self_adjoint_eigvals(x)
      #   psd_mask = tf.cast(
      #     tf.reduce_min(eigenvalues, axis=-1) >= 0, dtype=x.dtype)
      #   tf.where(psd_mask, answer, float('-inf'))
      # to emit probability 0 for inputs that are not PSD, without ever raising
      # an error. More care must be taken, as due to numerical stability issues,
      # self_adjoint_eigvals can return slightly negative eigenvalues even for
      # a PSD matrix.
      if self.input_output_cholesky:
        logdet = 2.0 * tf.reduce_sum(
            input_tensor=tf.math.log(tf.linalg.diag_part(x)), axis=[-1])
      else:
        _, logdet = tf.linalg.slogdet(x)
      answer = (self.concentration - 1.) * logdet
      return answer

Response: Returns the unnormalized log density of a LKJ distribution.


Instruction:
How would you code a function in Python 3 to
return the log normalization of an LKJ distribution.

Response:
def _log_normalization(self, name='log_normalization'):
    """Returns the log normalization of an LKJ distribution.

    Args:
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      log_z: A Tensor of the same shape and dtype as `concentration`, containing
        the corresponding log normalizers.
    """
    # The formula is from D. Lewandowski et al [1], p. 1999, from the
    # proof that eqs 16 and 17 are equivalent.
    with tf.name_scope(name or 'log_normalization_lkj'):
      logpi = np.log(np.pi)
      ans = tf.zeros_like(self.concentration)
      for k in range(1, self.dimension):
        ans += logpi * (k / 2.)
        ans += tf.math.lgamma(self.concentration +
                              (self.dimension - 1 - k) / 2.)
        ans -= tf.math.lgamma(self.concentration + (self.dimension - 1) / 2.)
      return ans


Instruction:
Given the following Python 3 function, write the documentation
def common_dtype(args_list, preferred_dtype=None):
  dtype = None
  preferred_dtype = (None if preferred_dtype is None
                     else tf.as_dtype(preferred_dtype))
  for a in tf.nest.flatten(args_list):
    if hasattr(a, 'dtype'):
      dt = tf.as_dtype(a.dtype)
    else:
      continue
    if dtype is None:
      dtype = dt
    elif dtype != dt:
      raise TypeError('Found incompatible dtypes, {} and {}.'.format(dtype, dt))
  if dtype is None and preferred_dtype is None:
    return None
  return (preferred_dtype if dtype is None else dtype).as_numpy_dtype

Response: Returns explict dtype from args_list if exists else preferred_dtype.


Instruction:
Given the following Python 3 function, write the documentation
def _make_summary_statistic(attr):
  def _fn(self, **kwargs):
    """Implements summary statistic, eg, mean, stddev, mode."""
    x = getattr(self.distribution, attr)(**kwargs)
    shape = prefer_static.concat([
        self.distribution.batch_shape_tensor(),
        prefer_static.ones(prefer_static.rank_from_shape(self.sample_shape),
                           dtype=self.sample_shape.dtype),
        self.distribution.event_shape_tensor(),
    ], axis=0)
    x = tf.reshape(x, shape=shape)
    shape = prefer_static.concat([
        self.distribution.batch_shape_tensor(),
        self.sample_shape,
        self.distribution.event_shape_tensor(),
    ], axis=0)
    return tf.broadcast_to(x, shape)
  return _fn

Response: Factory for implementing summary statistics eg mean stddev mode.


Instruction:
Can you generate a brief explanation for the following Python 3 code
def _kl_sample(a, b, name='kl_sample'):
  assertions = []
  a_ss = tf.get_static_value(a.sample_shape)
  b_ss = tf.get_static_value(b.sample_shape)
  msg = '`a.sample_shape` must be identical to `b.sample_shape`.'
  if a_ss is not None and b_ss is not None:
    if not np.array_equal(a_ss, b_ss):
      raise ValueError(msg)
  elif a.validate_args or b.validate_args:
    assertions.append(assert_util.assert_equal(
        a.sample_shape, b.sample_shape, message=msg))
  with tf.control_dependencies(assertions):
    kl = kullback_leibler.kl_divergence(
        a.distribution, b.distribution, name=name)
    n = prefer_static.reduce_prod(a.sample_shape)
    return tf.cast(x=n, dtype=kl.dtype) * kl

Response: Batchwise KL divergence of a and b.


Instruction:
Explain what the following Python 3 code does
def _broadcast_to(tensor_to_broadcast, target_tensors):
  output = tensor_to_broadcast
  for tensor in target_tensors:
    output += tf.zeros_like(tensor)
  return output

Response: Helper to broadcast a tensor to a list of target tensors.


Instruction:
Implement a Python 3 function for
evaluating the pdf at the peak.

Response:
def _pdf_at_peak(self):
    """Pdf evaluated at the peak."""
    return (self.peak - self.low) / (self.high - self.low)


Instruction:
Can you generate the documentation for the following Python 3 function
def effective_sample_size(states,
                          filter_threshold=0.,
                          filter_beyond_lag=None,
                          name=None):
  """Estimate a lower bound on effective sample size for each independent chain.

  Roughly speaking, "effective sample size" (ESS) is the size of an iid sample
  with the same variance as `state`.

  More precisely, given a stationary sequence of possibly correlated random
  variables `X_1, X_2,...,X_N`, each identically distributed ESS is the number
  such that

  ```Variance{ N**-1 * Sum{X_i} } = ESS**-1 * Variance{ X_1 }.```

  If the sequence is uncorrelated, `ESS = N`.  In general, one should expect
  `ESS <= N`, with more highly correlated sequences having smaller `ESS`.

  Args:
    states:  `Tensor` or list of `Tensor` objects.  Dimension zero should index
      identically distributed states.
    filter_threshold:  `Tensor` or list of `Tensor` objects.
      Must broadcast with `state`.  The auto-correlation sequence is truncated
      after the first appearance of a term less than `filter_threshold`.
      Setting to `None` means we use no threshold filter.  Since `|R_k| <= 1`,
      setting to any number less than `-1` has the same effect.
    filter_beyond_lag:  `Tensor` or list of `Tensor` objects.  Must be
      `int`-like and scalar valued.  The auto-correlation sequence is truncated
      to this length.  Setting to `None` means we do not filter based on number
      of lags.
    name:  `String` name to prepend to created ops.

  Returns:
    ess:  `Tensor` or list of `Tensor` objects.  The effective sample size of
      each component of `states`.  Shape will be `states.shape[1:]`.

  Raises:
    ValueError:  If `states` and `filter_threshold` or `states` and
      `filter_beyond_lag` are both lists with different lengths.

  #### Examples

  We use ESS to estimate standard error.

  ```
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  target = tfd.MultivariateNormalDiag(scale_diag=[1., 2.])

  # Get 1000 states from one chain.
  states = tfp.mcmc.sample_chain(
      num_burnin_steps=200,
      num_results=1000,
      current_state=tf.constant([0., 0.]),
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target.log_prob,
        step_size=0.05,
        num_leapfrog_steps=20))
  states.shape
  ==> (1000, 2)

  ess = effective_sample_size(states)
  ==> Shape (2,) Tensor

  mean, variance = tf.nn.moments(states, axis=0)
  standard_error = tf.sqrt(variance / ess)
  ```

  Some math shows that, with `R_k` the auto-correlation sequence,
  `R_k := Covariance{X_1, X_{1+k}} / Variance{X_1}`, we have

  ```ESS(N) =  N / [ 1 + 2 * ( (N - 1) / N * R_1 + ... + 1 / N * R_{N-1}  ) ]```

  This function estimates the above by first estimating the auto-correlation.
  Since `R_k` must be estimated using only `N - k` samples, it becomes
  progressively noisier for larger `k`.  For this reason, the summation over
  `R_k` should be truncated at some number `filter_beyond_lag < N`.  Since many
  MCMC methods generate chains where `R_k > 0`, a reasonable criteria is to
  truncate at the first index where the estimated auto-correlation becomes
  negative.

  The arguments `filter_beyond_lag`, `filter_threshold` are filters intended to
  remove noisy tail terms from `R_k`.  They combine in an "OR" manner meaning
  terms are removed if they were to be filtered under the `filter_beyond_lag` OR
  `filter_threshold` criteria.
  """
  states_was_list = _is_list_like(states)

  # Convert all args to lists.
  if not states_was_list:
    states = [states]

  filter_beyond_lag = _broadcast_maybelist_arg(states, filter_beyond_lag,
                                               'filter_beyond_lag')
  filter_threshold = _broadcast_maybelist_arg(states, filter_threshold,
                                              'filter_threshold')

  # Process items, one at a time.
  with tf.compat.v1.name_scope(name, 'effective_sample_size'):
    ess_list = [
        _effective_sample_size_single_state(s, ml, mlt)
        for (s, ml, mlt) in zip(states, filter_beyond_lag, filter_threshold)
    ]

  if states_was_list:
    return ess_list
  return ess_list[0]

Response: Estimate a lower bound on effective sample size for each independent chain.


Instruction:
Can you tell what is the following Python 3 function doing
def potential_scale_reduction(chains_states,
                              independent_chain_ndims=1,
                              name=None):
  """Gelman and Rubin (1992)'s potential scale reduction for chain convergence.

  Given `N > 1` states from each of `C > 1` independent chains, the potential
  scale reduction factor, commonly referred to as R-hat, measures convergence of
  the chains (to the same target) by testing for equality of means.
  Specifically, R-hat measures the degree to which variance (of the means)
  between chains exceeds what one would expect if the chains were identically
  distributed. See [Gelman and Rubin (1992)][1]; [Brooks and Gelman (1998)][2].

  Some guidelines:

  * The initial state of the chains should be drawn from a distribution
    overdispersed with respect to the target.
  * If all chains converge to the target, then as `N --> infinity`, R-hat --> 1.
    Before that, R-hat > 1 (except in pathological cases, e.g. if the chain
    paths were identical).
  * The above holds for any number of chains `C > 1`.  Increasing `C` does
    improves effectiveness of the diagnostic.
  * Sometimes, R-hat < 1.2 is used to indicate approximate convergence, but of
    course this is problem dependent. See [Brooks and Gelman (1998)][2].
  * R-hat only measures non-convergence of the mean. If higher moments, or
    other statistics are desired, a different diagnostic should be used. See
    [Brooks and Gelman (1998)][2].

  Args:
    chains_states:  `Tensor` or Python `list` of `Tensor`s representing the
      state(s) of a Markov Chain at each result step.  The `ith` state is
      assumed to have shape `[Ni, Ci1, Ci2,...,CiD] + A`.
      Dimension `0` indexes the `Ni > 1` result steps of the Markov Chain.
      Dimensions `1` through `D` index the `Ci1 x ... x CiD` independent
      chains to be tested for convergence to the same target.
      The remaining dimensions, `A`, can have any shape (even empty).
    independent_chain_ndims: Integer type `Tensor` with value `>= 1` giving the
      number of giving the number of dimensions, from `dim = 1` to `dim = D`,
      holding independent chain results to be tested for convergence.
    name: `String` name to prepend to created tf.  Default:
      `potential_scale_reduction`.

  Returns:
    `Tensor` or Python `list` of `Tensor`s representing the R-hat statistic for
    the state(s).  Same `dtype` as `state`, and shape equal to
    `state.shape[1 + independent_chain_ndims:]`.

  Raises:
    ValueError:  If `independent_chain_ndims < 1`.

  #### Examples

  Diagnosing convergence by monitoring 10 chains that each attempt to
  sample from a 2-variate normal.

  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  target = tfd.MultivariateNormalDiag(scale_diag=[1., 2.])

  # Get 10 (2x) overdispersed initial states.
  initial_state = target.sample(10) * 2.
  ==> (10, 2)

  # Get 1000 samples from the 10 independent chains.
  chains_states, _ = tfp.mcmc.sample_chain(
      num_burnin_steps=200,
      num_results=1000,
      current_state=initial_state,
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target.log_prob,
          step_size=0.05,
          num_leapfrog_steps=20))
  chains_states.shape
  ==> (1000, 10, 2)

  rhat = tfp.mcmc.diagnostic.potential_scale_reduction(
      chains_states, independent_chain_ndims=1)

  # The second dimension needed a longer burn-in.
  rhat.eval()
  ==> [1.05, 1.3]
  ```

  To see why R-hat is reasonable, let `X` be a random variable drawn uniformly
  from the combined states (combined over all chains).  Then, in the limit
  `N, C --> infinity`, with `E`, `Var` denoting expectation and variance,

  ```R-hat = ( E[Var[X | chain]] + Var[E[X | chain]] ) / E[Var[X | chain]].```

  Using the law of total variance, the numerator is the variance of the combined
  states, and the denominator is the total variance minus the variance of the
  the individual chain means.  If the chains are all drawing from the same
  distribution, they will have the same mean, and thus the ratio should be one.

  #### References

  [1]: Stephen P. Brooks and Andrew Gelman. General Methods for Monitoring
       Convergence of Iterative Simulations. _Journal of Computational and
       Graphical Statistics_, 7(4), 1998.

  [2]: Andrew Gelman and Donald B. Rubin. Inference from Iterative Simulation
       Using Multiple Sequences. _Statistical Science_, 7(4):457-472, 1992.
  """
  chains_states_was_list = _is_list_like(chains_states)
  if not chains_states_was_list:
    chains_states = [chains_states]

  # tf.get_static_value returns None iff a constant value (as a numpy
  # array) is not efficiently computable.  Therefore, we try constant_value then
  # check for None.
  icn_const_ = tf.get_static_value(
      tf.convert_to_tensor(value=independent_chain_ndims))
  if icn_const_ is not None:
    independent_chain_ndims = icn_const_
    if icn_const_ < 1:
      raise ValueError(
          'Argument `independent_chain_ndims` must be `>= 1`, found: {}'.format(
              independent_chain_ndims))

  with tf.compat.v1.name_scope(name, 'potential_scale_reduction'):
    rhat_list = [
        _potential_scale_reduction_single_state(s, independent_chain_ndims)
        for s in chains_states
    ]

  if chains_states_was_list:
    return rhat_list
  return rhat_list[0]

Response: Gelman and Rubin 2002 s potential scale reduction for chains.


Instruction:
Can you generate a brief explanation for the following Python 3 code
def _potential_scale_reduction_single_state(state, independent_chain_ndims):
  with tf.compat.v1.name_scope(
      'potential_scale_reduction_single_state',
      values=[state, independent_chain_ndims]):
    # We assume exactly one leading dimension indexes e.g. correlated samples
    # from each Markov chain.
    state = tf.convert_to_tensor(value=state, name='state')
    sample_ndims = 1

    sample_axis = tf.range(0, sample_ndims)
    chain_axis = tf.range(sample_ndims,
                          sample_ndims + independent_chain_ndims)
    sample_and_chain_axis = tf.range(
        0, sample_ndims + independent_chain_ndims)

    n = _axis_size(state, sample_axis)
    m = _axis_size(state, chain_axis)

    # In the language of Brooks and Gelman (1998),
    # B / n is the between chain variance, the variance of the chain means.
    # W is the within sequence variance, the mean of the chain variances.
    b_div_n = _reduce_variance(
        tf.reduce_mean(input_tensor=state, axis=sample_axis, keepdims=True),
        sample_and_chain_axis,
        biased=False)
    w = tf.reduce_mean(
        input_tensor=_reduce_variance(
            state, sample_axis, keepdims=True, biased=True),
        axis=sample_and_chain_axis)

    # sigma^2_+ is an estimate of the true variance, which would be unbiased if
    # each chain was drawn from the target.  c.f. "law of total variance."
    sigma_2_plus = w + b_div_n

    return ((m + 1.) / m) * sigma_2_plus / w - (n - 1.) / (m * n)

Response: Internal function for the potential scale reduction for one single state.


Instruction:
Can you create a Python 3 function that
gets number of elements of x in axis as type x. dtype.

Response:
def _axis_size(x, axis=None):
  """Get number of elements of `x` in `axis`, as type `x.dtype`."""
  if axis is None:
    return tf.cast(tf.size(input=x), x.dtype)
  return tf.cast(
      tf.reduce_prod(input_tensor=tf.gather(tf.shape(input=x), axis)), x.dtype)


Instruction:
Can you generate a brief explanation for the following Python 3 code
def _broadcast_maybelist_arg(states, secondary_arg, name):
  if _is_list_like(secondary_arg):
    if len(secondary_arg) != len(states):
      raise ValueError('Argument `%s` was a list of different length ({}) than '
                       '`states` ({})'.format(name, len(states)))
  else:
    secondary_arg = [secondary_arg] * len(states)

  return secondary_arg

Response: Broadcast a listable secondary_arg to that of states.


Instruction:
Can you tell what is the following Python 3 function doing
def quadrature_scheme_lognormal_gauss_hermite(
    loc, scale, quadrature_size,
    validate_args=False, name=None):  # pylint: disable=unused-argument
  """Use Gauss-Hermite quadrature to form quadrature on positive-reals.

  Note: for a given `quadrature_size`, this method is generally less accurate
  than `quadrature_scheme_lognormal_quantiles`.

  Args:
    loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
      the LogNormal prior.
    scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
      the LogNormal prior.
    quadrature_size: Python `int` scalar representing the number of quadrature
      points.
    validate_args: Python `bool`, default `False`. When `True` distribution
      parameters are checked for validity despite possibly degrading runtime
      performance. When `False` invalid inputs may silently render incorrect
      outputs.
    name: Python `str` name prefixed to Ops created by this class.

  Returns:
    grid: (Batch of) length-`quadrature_size` vectors representing the
      `log_rate` parameters of a `Poisson`.
    probs: (Batch of) length-`quadrature_size` vectors representing the
      weight associate with each `grid` value.
  """
  with tf.name_scope(
      name or "vector_diffeomixture_quadrature_gauss_hermite"):
    grid, probs = np.polynomial.hermite.hermgauss(deg=quadrature_size)
    npdt = dtype_util.as_numpy_dtype(loc.dtype)
    grid = grid.astype(npdt)
    probs = probs.astype(npdt)
    probs /= np.linalg.norm(probs, ord=1, keepdims=True)
    probs = tf.convert_to_tensor(value=probs, name="probs", dtype=loc.dtype)
    # The following maps the broadcast of `loc` and `scale` to each grid
    # point, i.e., we are creating several log-rates that correspond to the
    # different Gauss-Hermite quadrature points and (possible) batches of
    # `loc` and `scale`.
    grid = (loc[..., tf.newaxis] + np.sqrt(2.) * scale[..., tf.newaxis] * grid)
    return grid, probs

Response: Use Gauss - Hermite quadrature to form a lognormal graph on positive - reals.


Instruction:
How would you implement a function in Python 3 that
uses LogNormal quantiles to form quadrature on positive - reals.

Response:
def quadrature_scheme_lognormal_quantiles(
    loc, scale, quadrature_size,
    validate_args=False, name=None):
  """Use LogNormal quantiles to form quadrature on positive-reals.

  Args:
    loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
      the LogNormal prior.
    scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
      the LogNormal prior.
    quadrature_size: Python `int` scalar representing the number of quadrature
      points.
    validate_args: Python `bool`, default `False`. When `True` distribution
      parameters are checked for validity despite possibly degrading runtime
      performance. When `False` invalid inputs may silently render incorrect
      outputs.
    name: Python `str` name prefixed to Ops created by this class.

  Returns:
    grid: (Batch of) length-`quadrature_size` vectors representing the
      `log_rate` parameters of a `Poisson`.
    probs: (Batch of) length-`quadrature_size` vectors representing the
      weight associate with each `grid` value.
  """
  with tf.name_scope(name or "quadrature_scheme_lognormal_quantiles"):
    # Create a LogNormal distribution.
    dist = transformed_distribution.TransformedDistribution(
        distribution=normal.Normal(loc=loc, scale=scale),
        bijector=exp_bijector.Exp(),
        validate_args=validate_args)
    batch_ndims = tensorshape_util.rank(dist.batch_shape)
    if batch_ndims is None:
      batch_ndims = tf.shape(input=dist.batch_shape_tensor())[0]

    def _compute_quantiles():
      """Helper to build quantiles."""
      # Omit {0, 1} since they might lead to Inf/NaN.
      zero = tf.zeros([], dtype=dist.dtype)
      edges = tf.linspace(zero, 1., quadrature_size + 3)[1:-1]
      # Expand edges so its broadcast across batch dims.
      edges = tf.reshape(
          edges,
          shape=tf.concat(
              [[-1], tf.ones([batch_ndims], dtype=tf.int32)], axis=0))
      quantiles = dist.quantile(edges)
      # Cyclically permute left by one.
      perm = tf.concat([tf.range(1, 1 + batch_ndims), [0]], axis=0)
      quantiles = tf.transpose(a=quantiles, perm=perm)
      return quantiles
    quantiles = _compute_quantiles()

    # Compute grid as quantile midpoints.
    grid = (quantiles[..., :-1] + quantiles[..., 1:]) / 2.
    # Set shape hints.
    new_shape = tensorshape_util.concatenate(dist.batch_shape,
                                             [quadrature_size])
    tensorshape_util.set_shape(grid, new_shape)

    # By construction probs is constant, i.e., `1 / quadrature_size`. This is
    # important, because non-constant probs leads to non-reparameterizable
    # samples.
    probs = tf.fill(
        dims=[quadrature_size], value=1. / tf.cast(quadrature_size, dist.dtype))

    return grid, probs


Instruction:
Can you generate a brief explanation for the following Python 3 code
def merge(self, x=None, y=None, ildj=None, kwargs=None, mapping=None):
    if mapping is None:
      mapping = _Mapping(x=x, y=y, ildj=ildj, kwargs=kwargs)
    elif any(arg is not None for arg in [x, y, ildj, kwargs]):
      raise ValueError("Cannot simultaneously specify mapping and individual "
                       "arguments.")

    return _Mapping(
        x=self._merge(self.x, mapping.x),
        y=self._merge(self.y, mapping.y),
        ildj=self._merge(self.ildj, mapping.ildj),
        kwargs=self._merge(self.kwargs, mapping.kwargs, use_equals=True))

Response: Returns a new _Mapping with args merged with self.


Instruction:
Implement a function in Python 3 to
remove the given field from the cache value.

Response:
def remove(self, field):
    """To support weak referencing, removes cache key from the cache value."""
    return _Mapping(
        x=None if field == "x" else self.x,
        y=None if field == "y" else self.y,
        ildj=self.ildj,
        kwargs=self.kwargs)


Instruction:
Here you have a function in Python 3, explain what it does
def _merge(self, old, new, use_equals=False):
    if old is None:
      return new
    if new is None:
      return old
    if (old == new) if use_equals else (old is new):
      return old
    raise ValueError("Incompatible values: %s != %s" % (old, new))

Response: Helper to merge two values.


Instruction:
Implement a Python 3 function for
converting nested tuple list or dict to nested tuple.

Response:
def _deep_tuple(self, x):
    """Converts nested `tuple`, `list`, or `dict` to nested `tuple`."""
    if isinstance(x, dict):
      return self._deep_tuple(tuple(sorted(x.items())))
    elif isinstance(x, (list, tuple)):
      return tuple(map(self._deep_tuple, x))

    return x


Instruction:
Explain what the following Python 3 code does
def _left_doubling_increments(batch_shape, max_doublings, step_size, seed=None,
                              name=None):
  """Computes the doubling increments for the left end point.

  The doubling procedure expands an initial interval to find a superset of the
  true slice. At each doubling iteration, the interval width is doubled to
  either the left or the right hand side with equal probability.
  If, initially, the left end point is at `L(0)` and the width of the
  interval is `w(0)`, then the left end point and the width at the
  k-th iteration (denoted L(k) and w(k) respectively) are given by the following
  recursions:

  ```none
  w(k) = 2 * w(k-1)
  L(k) = L(k-1) - w(k-1) * X_k, X_k ~ Bernoulli(0.5)
  or, L(0) - L(k) = w(0) Sum(2^i * X(i+1), 0 <= i < k)
  ```

  This function computes the sequence of `L(0)-L(k)` and `w(k)` for k between 0
  and `max_doublings` independently for each chain.

  Args:
    batch_shape: Positive int32 `tf.Tensor`. The batch shape.
    max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
      doublings to consider.
    step_size: A real `tf.Tensor` with shape compatible with [num_chains].
      The size of the initial interval.
    seed: (Optional) positive int. The random seed. If None, no seed is set.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    left_increments: A tensor of shape (max_doublings+1, batch_shape). The
      relative position of the left end point after the doublings.
    widths: A tensor of shape (max_doublings+1, ones_like(batch_shape)). The
      widths of the intervals at each stage of the doubling.
  """
  with tf.compat.v1.name_scope(name, 'left_doubling_increments',
                               [batch_shape, max_doublings, step_size]):

    step_size = tf.convert_to_tensor(value=step_size)
    dtype = step_size.dtype.base_dtype
    # Output shape of the left increments tensor.
    output_shape = tf.concat(([max_doublings + 1], batch_shape), axis=0)
    # A sample realization of X_k.
    expand_left = distributions.Bernoulli(0.5, dtype=dtype).sample(
        sample_shape=output_shape, seed=seed)

    # The widths of the successive intervals. Starts with 1.0 and ends with
    # 2^max_doublings.
    width_multipliers = tf.cast(2 ** tf.range(0, max_doublings+1), dtype=dtype)
    # Output shape of the `widths` tensor.
    widths_shape = tf.concat(([max_doublings + 1],
                              tf.ones_like(batch_shape)), axis=0)
    width_multipliers = tf.reshape(width_multipliers, shape=widths_shape)
    # Widths shape is [max_doublings + 1, 1, 1, 1...].
    widths = width_multipliers * step_size

    # Take the cumulative sum of the left side increments in slice width to give
    # the resulting distance from the inital lower bound.
    left_increments = tf.cumsum(widths * expand_left, exclusive=True, axis=0)
    return left_increments, widths

Response: This function computes the doubling increments for the left end point.


Instruction:
Explain what the following Python 3 code does
def _find_best_interval_idx(x, name=None):
  with tf.compat.v1.name_scope(name, 'find_best_interval_idx', [x]):
    # Returns max_doublings + 1. Positive int32.
    k = tf.shape(input=x)[0]
    dtype = x.dtype.base_dtype
    # Factors by which to multiply the flag. Corresponds to (2 * k - i) above.
    mults = tf.range(2 * k, k, -1, dtype=dtype)[:, tf.newaxis]
    # Factors by which to shift the flag. Corresponds to i above. Ensures the
    # widest bounds are selected if there are no bounds outside the slice.
    shifts = tf.range(k, dtype=dtype)[:, tf.newaxis]
    indices = tf.argmax(input=mults * x + shifts, axis=0, output_type=dtype)
    return indices

Response: Internal function for finding the index of the optimal set of bounds for each chain.


Instruction:
Can you tell what is the following Python 3 function doing
def slice_bounds_by_doubling(x_initial,
                             target_log_prob,
                             log_slice_heights,
                             max_doublings,
                             step_size,
                             seed=None,
                             name=None):
  """Returns the bounds of the slice at each stage of doubling procedure.

  Precomputes the x coordinates of the left (L) and right (R) endpoints of the
  interval `I` produced in the "doubling" algorithm [Neal 2003][1] P713. Note
  that we simultaneously compute all possible doubling values for each chain,
  for the reason that at small-medium densities, the gains from parallel
  evaluation might cause a speed-up, but this will be benchmarked against the
  while loop implementation.

  Args:
    x_initial: `tf.Tensor` of any shape and any real dtype consumable by
      `target_log_prob`. The initial points.
    target_log_prob: A callable taking a `tf.Tensor` of shape and dtype as
      `x_initial` and returning a tensor of the same shape. The log density of
      the target distribution.
    log_slice_heights: `tf.Tensor` with the same shape as `x_initial` and the
      same dtype as returned by `target_log_prob`. The log of the height of the
      slice for each chain. The values must be bounded above by
      `target_log_prob(x_initial)`.
    max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
      doublings to consider.
    step_size: `tf.Tensor` with same dtype as and shape compatible with
      `x_initial`. The size of the initial interval.
    seed: (Optional) positive int. The random seed. If None, no seed is set.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    upper_bounds: A tensor of same shape and dtype as `x_initial`. Slice upper
      bounds for each chain.
    lower_bounds: A tensor of same shape and dtype as `x_initial`. Slice lower
      bounds for each chain.
    both_ok: A tensor of shape `x_initial` and boolean dtype. Indicates if both
      the chosen upper and lower bound lie outside of the slice.

  #### References

  [1]: Radford M. Neal. Slice Sampling. The Annals of Statistics. 2003, Vol 31,
       No. 3 , 705-767.
       https://projecteuclid.org/download/pdf_1/euclid.aos/1056562461
  """
  with tf.compat.v1.name_scope(
      name, 'slice_bounds_by_doubling',
      [x_initial, log_slice_heights, max_doublings, step_size]):
    seed_gen = distributions.SeedStream(seed, salt='slice_bounds_by_doubling')
    x_initial = tf.convert_to_tensor(value=x_initial)
    batch_shape = tf.shape(input=x_initial)
    dtype = step_size.dtype.base_dtype
    left_endpoints = x_initial + step_size * tf.random.uniform(
        batch_shape, minval=-1.0, maxval=0.0, dtype=dtype, seed=seed_gen())

    # Compute the increments by which we need to step the upper and lower bounds
    # part of the doubling procedure.
    left_increments, widths = _left_doubling_increments(
        batch_shape, max_doublings, step_size, seed=seed_gen())
    # The left and right end points. Shape (max_doublings+1,) + batch_shape.
    left_endpoints -= left_increments
    right_endpoints = left_endpoints + widths

    # Test if these end points lie outside of the slice.
    # Checks if the end points of the slice are outside the graph of the pdf.
    left_ep_values = tf.map_fn(target_log_prob, left_endpoints)
    right_ep_values = tf.map_fn(target_log_prob, right_endpoints)
    left_ok = left_ep_values < log_slice_heights
    right_ok = right_ep_values < log_slice_heights
    both_ok = left_ok & right_ok

    both_ok_f = tf.reshape(both_ok, [max_doublings + 1, -1])

    best_interval_idx = _find_best_interval_idx(
        tf.cast(both_ok_f, dtype=tf.int32))

    # Formats the above index as required to use with gather_nd.
    point_index_gather = tf.stack(
        [best_interval_idx,
         tf.range(tf.size(input=best_interval_idx))],
        axis=1,
        name='point_index_gather')
    left_ep_f = tf.reshape(left_endpoints, [max_doublings + 1, -1])
    right_ep_f = tf.reshape(right_endpoints, [max_doublings + 1, -1])
    # The x values of the uppper and lower bounds of the slices for each chain.
    lower_bounds = tf.reshape(tf.gather_nd(left_ep_f, point_index_gather),
                              batch_shape)
    upper_bounds = tf.reshape(tf.gather_nd(right_ep_f, point_index_gather),
                              batch_shape)
    both_ok = tf.reduce_any(input_tensor=both_ok, axis=0)
    return upper_bounds, lower_bounds, both_ok

Response: Returns the upper and lower bounds of a slice at each stage of doubling procedure.


Instruction:
Here you have a function in Python 3, explain what it does
def _sample_with_shrinkage(x_initial, target_log_prob, log_slice_heights,
                           step_size, lower_bounds, upper_bounds, seed=None,
                           name=None):
  """Samples from the slice by applying shrinkage for rejected points.

  Implements the one dimensional slice sampling algorithm of Neal (2003), with a
  doubling algorithm (Neal 2003 P715 Fig. 4), which doubles the size of the
  interval at each iteration and shrinkage (Neal 2003 P716 Fig. 5), which
  reduces the width of the slice when a selected point is rejected, by setting
  the relevant bound that that value. Randomly sampled points are checked for
  two criteria: that they lie within the slice and that they pass the
  acceptability check (Neal 2003 P717 Fig. 6), which tests that the new state
  could have generated the previous one.

  Args:
    x_initial: A tensor of any shape. The initial positions of the chains. This
      function assumes that all the dimensions of `x_initial` are batch
      dimensions (i.e. the event shape is `[]`).
    target_log_prob: Callable accepting a tensor like `x_initial` and returning
      a tensor containing the log density at that point of the same shape.
    log_slice_heights: Tensor of the same shape and dtype as the return value
      of `target_log_prob` when applied to `x_initial`. The log of the height of
      the chosen slice.
    step_size: A tensor of shape and dtype compatible with `x_initial`. The min
      interval size in the doubling algorithm.
    lower_bounds: Tensor of same shape and dtype as `x_initial`. Slice lower
      bounds for each chain.
    upper_bounds: Tensor of same shape and dtype as `x_initial`. Slice upper
      bounds for each chain.
    seed: (Optional) positive int. The random seed. If None, no seed is set.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    x_proposed: A tensor of the same shape and dtype as `x_initial`. The next
      proposed state of the chain.
  """
  with tf.compat.v1.name_scope(
      name, 'sample_with_shrinkage',
      [x_initial, log_slice_heights, step_size, lower_bounds, upper_bounds]):
    seed_gen = distributions.SeedStream(seed, salt='_sample_with_shrinkage')
    # Keeps track of whether an acceptable sample has been found for the chain.
    found = tf.zeros_like(x_initial, dtype=tf.bool)
    cond = lambda found, *ignored_args: ~tf.reduce_all(input_tensor=found)
    x_next = tf.identity(x_initial)
    x_initial_shape = tf.shape(input=x_initial)
    x_initial_dtype = x_initial.dtype.base_dtype
    def _body(found, left, right, x_next):
      """Iterates until every chain has found a suitable next state."""
      proportions = tf.random.uniform(
          x_initial_shape, dtype=x_initial_dtype, seed=seed_gen())
      x_proposed = tf.where(~found, left + proportions * (right - left), x_next)
      accept_res = _test_acceptance(x_initial, target_log_prob=target_log_prob,
                                    decided=found,
                                    log_slice_heights=log_slice_heights,
                                    x_proposed=x_proposed, step_size=step_size,
                                    lower_bounds=left, upper_bounds=right)
      boundary_test = log_slice_heights < target_log_prob(x_proposed)
      can_accept = boundary_test & accept_res
      next_found = found | can_accept
      # Note that it might seem that we are moving the left and right end points
      # even if the point has been accepted (which is contrary to the stated
      # algorithm in Neal). However, this does not matter because the endpoints
      # for points that have been already accepted are not used again so it
      # doesn't matter what we do with them.
      next_left = tf.where(x_proposed < x_initial, x_proposed, left)
      next_right = tf.where(x_proposed >= x_initial, x_proposed, right)
      return next_found, next_left, next_right, x_proposed

    return tf.while_loop(
        cond=cond,
        body=_body,
        loop_vars=(found, lower_bounds, upper_bounds, x_next))[-1]

Response: This function is a helper function for the one dimensional slice sampling algorithm of Neal 2003.