Instruction:
Can you generate the documentation for the following Python 3 function
def minimize(objective_function,
             initial_population=None,
             initial_position=None,
             population_size=50,
             population_stddev=1.,
             max_iterations=100,
             func_tolerance=0,
             position_tolerance=1e-8,
             differential_weight=0.5,
             crossover_prob=0.9,
             seed=None,
             name=None):
  """Applies the Differential evolution algorithm to minimize a function.

  Differential Evolution is an evolutionary optimization algorithm which works
  on a set of candidate solutions called the population. It iteratively
  improves the population by applying genetic operators of mutation and
  recombination. The objective function `f` supplies the fitness of each
  candidate. A candidate `s_1` is considered better than `s_2` if
  `f(s_1) < f(s_2)`.

  This method allows the user to either specify an initial population or a
  single candidate solution. If a single solution is specified, a population
  of the specified size is initialized by adding independent normal noise
  to the candidate solution.

  The implementation also supports a multi-part specification of the state. For
  example, consider the objective function:

  ```python
  # x is a tensor of shape [n, m] while y is of shape [n].
  def objective(x, y):
    return tf.math.reduce_sum(x ** 2, axis=-1) + y ** 2
  ```
  The state in this case is specified by two input tensors `x` and `y`. To
  apply the algorithm to this objective function, one would need to specify
  either an initial population as a list of two tensors of shapes
  `[population_size, k]` and `[population_size]`. The following code shows the
  complete example:

  ```python
    population_size = 40
    # With an initial population and a multi-part state.
    initial_population = (tf.random.normal([population_size]),
                          tf.random.normal([population_size]))
    def easom_fn(x, y):
      return -(tf.math.cos(x) * tf.math.cos(y) *
               tf.math.exp(-(x-np.pi)**2 - (y-np.pi)**2))

    optim_results = tfp.optimizers.differential_evolution_minimize(
        easom_fn,
        initial_population=initial_population,
        seed=43210)

    print (optim_results.converged)
    print (optim_results.position)  # Should be (close to) [pi, pi].
    print (optim_results.objective_value)    # Should be -1.

    # With a single starting point
    initial_position = (tf.constant(1.0), tf.constant(1.0))

    optim_results = tfp.optimizers.differential_evolution_minimize(
        easom_fn,
        initial_position=initial_position,
        population_size=40,
        population_stddev=2.0,
        seed=43210)
  ```

  Args:
    objective_function: A Python callable that accepts a batch of possible
      solutions and returns the values of the objective function at those
      arguments as a rank 1 real `Tensor`. This specifies the function to be
      minimized. The input to this callable may be either a single `Tensor`
      or a Python `list` of `Tensor`s. The signature must match the format of
      the argument `population`. (i.e. objective_function(*population) must
      return the value of the function to be minimized).
    initial_population: A real `Tensor` or Python list of `Tensor`s.
      If a list, each `Tensor` must be of rank at least 1 and with a common
      first dimension. The first dimension indexes into the candidate solutions
      while the rest of the dimensions (if any) index into an individual
      solution. The size of the population must be at least 4. This is a
      requirement of the DE algorithm.
    initial_position: A real `Tensor` of any shape. The seed solution used
      to initialize the population of solutions. If this parameter is specified
      then `initial_population` must not be specified.
    population_size: A positive scalar int32 `Tensor` greater than 4. The
      size of the population to evolve. This parameter is ignored if
      `initial_population` is specified.
      Default value: 50.
    population_stddev: A positive scalar real `Tensor` of the same dtype
      as `initial_position`. This parameter is ignored if `initial_population`
      is specified. Used to generate the population from the `initial_position`
      by adding random normal noise with zero mean and the specified standard
      deviation.
      Default value: 1.0
    max_iterations: Positive scalar int32 `Tensor`. The maximum number of
      generations to evolve the population for.
      Default value: 100
    func_tolerance: Scalar `Tensor` of the same dtype as the output of the
      `objective_function`. The algorithm stops if the absolute difference
      between the largest and the smallest objective function value in the
      population is below this number.
      Default value: 0
    position_tolerance: Scalar `Tensor` of the same real dtype as
      `initial_position` or `initial_population`. The algorithm terminates if
      the largest absolute difference between the coordinates of the population
      members is below this threshold.
      Default value: 1e-8
    differential_weight: Real scalar `Tensor`. Must be positive and less than
      2.0. The parameter controlling the strength of mutation in the algorithm.
      Default value: 0.5
    crossover_prob: Real scalar `Tensor`. Must be between 0 and 1. The
      probability of recombination per site.
      Default value: 0.9
    seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
      applied.
      Default value: None.
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name
      'differential_evolution_minimize' is used.
      Default value: None

  Returns:
    optimizer_results: An object containing the following attributes:
      converged: Scalar boolean `Tensor` indicating whether the minimum was
        found within the specified tolerances.
      num_objective_evaluations: The total number of objective
        evaluations performed.
      position: A `Tensor` containing the best point found during the search.
        If the search converged, then this value is the argmin of the
        objective function within the specified tolerances.
      objective_value: A `Tensor` containing the value of the objective
        function at the `position`. If the search
        converged, then this is the (local) minimum of
        the objective function.
      final_population: The final state of the population.
      final_objective_values: The objective function evaluated at the
        final population.
      initial_population: The starting population.
      initial_objective_values: The objective function evaluated at the
        initial population.
      num_iterations: The number of iterations of the main algorithm body.

  Raises:
    ValueError: If neither the initial population, nor the initial position
      are specified or if both are specified.
  """

  if initial_population is None and initial_position is None:
    raise ValueError('Either the initial population or the initial position '
                     'must be specified.')
  if initial_population is not None and initial_position is not None:
    raise ValueError('Only one of initial population or initial position '
                     'should be specified')

  with tf.compat.v1.name_scope(
      name,
      default_name='minimize',
      values=[
          initial_population, initial_position, population_size,
          population_stddev, max_iterations, func_tolerance, position_tolerance,
          differential_weight, crossover_prob
      ]):
    (
        was_iterable,
        population,
        population_values,
        max_iterations,
        func_tolerance,
        position_tolerance,
        differential_weight,
        crossover_prob
    ) = _get_initial_args(objective_function,
                          initial_population,
                          initial_position,
                          population_size,
                          population_stddev,
                          max_iterations,
                          func_tolerance,
                          position_tolerance,
                          differential_weight,
                          crossover_prob,
                          seed)

    def evolve_body(loop_vars):
      """Performs one step of the evolution."""
      next_population, next_population_values = one_step(
          objective_function,
          loop_vars.population,
          population_values=loop_vars.population_values,
          differential_weight=differential_weight,
          crossover_prob=crossover_prob,
          seed=seed)
      converged = _check_convergence(next_population,
                                     next_population_values,
                                     func_tolerance,
                                     position_tolerance)

      failed = _check_failure(next_population_values)

      return [_MinimizeLoopVars(
          converged=converged,
          failed=failed,
          num_iterations=loop_vars.num_iterations+1,
          population=next_population,
          population_values=next_population_values)]

    def evolve_cond(loop_vars):
      should_stop = (
          loop_vars.failed |
          loop_vars.converged |
          (max_iterations is not None and
           loop_vars.num_iterations >= max_iterations))
      return ~should_stop

    initial_vars = _MinimizeLoopVars(
        converged=tf.convert_to_tensor(value=False),
        failed=tf.convert_to_tensor(value=False),
        num_iterations=tf.convert_to_tensor(value=0),
        population=population,
        population_values=population_values)
    final_state = tf.while_loop(
        cond=evolve_cond, body=evolve_body, loop_vars=(initial_vars,))[0]
    best_position, best_values = _find_best_in_population(
        final_state.population,
        final_state.population_values)
    # Ensure we return a similar structure to what the user supplied.
    final_population = final_state.population
    if not was_iterable:
      final_population = final_population[0]
      best_position = best_position[0]
    return DifferentialEvolutionOptimizerResults(
        converged=final_state.converged,
        failed=final_state.failed,
        position=best_position,
        objective_value=best_values,
        final_population=final_population,
        final_objective_values=final_state.population_values,
        initial_population=population,
        initial_objective_values=population_values,
        num_iterations=final_state.num_iterations)

Response: Applies the objective function to minimize a function.


Instruction:
Can you write a function in Python 3 where it
checks if all the population values are NaN or infinite.

Response:
def _check_failure(population_values):
  """Checks if all the population values are NaN/infinite."""
  return tf.math.reduce_all(input_tensor=tf.math.is_inf(population_values))


Instruction:
Can you create a Python 3 function that
finds the population member with the lowest value.

Response:
def _find_best_in_population(population, values):
  """Finds the population member with the lowest value."""
  best_value = tf.math.reduce_min(input_tensor=values)
  best_index = tf.where(tf.math.equal(values, best_value))[0, 0]

  return ([population_part[best_index] for population_part in population],
          best_value)


Instruction:
Can you generate a brief explanation for the following Python 3 code
def _check_convergence(population,
                       population_values,
                       func_tolerance,
                       position_tolerance):
  """Checks whether the convergence criteria have been met."""
  # Check func tolerance
  value_range = tf.math.abs(
      tf.math.reduce_max(input_tensor=population_values) -
      tf.math.reduce_min(input_tensor=population_values))
  value_converged = value_range <= func_tolerance
  # Ideally, we would compute the position convergence by computing the
  # pairwise distance between every member of the population and checking if
  # the maximum of those is less than the supplied tolerance. However, this is
  # completely infeasible in terms of performance. We adopt a more conservative
  # approach which checks the distance between the first population member
  # with the rest of the population. If the largest such distance is less than
  # half the supplied tolerance, we stop. The reason why this is sufficient is
  # as follows. For any pair of distinct points (a, b) in the population, we
  # have the relation:  |a - b| <= |x0 - a| + |x0 - b|, where x0 is any
  # other point. In particular, let x0 be the first element of the population
  # and suppose that the largest distance between this point and any other
  # member is epsilon. Then, for any pair of points (a, b),
  # |a - b| <= 2 * epsilon and hence, the maximum distance between any pair of
  # points in the population is bounded above by twice the distance between
  # the first point and other points.
  half_tol = position_tolerance / 2
  def part_converged(part):
    return tf.math.reduce_max(input_tensor=tf.math.abs(part -
                                                       part[0])) <= half_tol

  x_converged = tf.math.reduce_all(
      input_tensor=[part_converged(part) for part in population])
  return value_converged | x_converged

Response: Checks whether the convergence criteria have been met.


Instruction:
Here you have a function in Python 3, explain what it does
def _get_starting_population(initial_population,
                             initial_position,
                             population_size,
                             population_stddev,
                             seed):
  """Constructs the initial population.

  If an initial population is not already provided, this function constructs
  a population by adding random normal noise to the initial position.

  Args:
    initial_population: None or a list of `Tensor`s. The initial population.
    initial_position: None or a list of `Tensor`s. The initial position.
      If initial_population is None, this argument must not be None.
    population_size: Scalar integer `Tensor`. The number of members in the
      population. If the initial population is not None, this parameter is
      ignored.
    population_stddev: A positive scalar real `Tensor` of the same dtype
      as `initial_position` or `initial_population` (whichever is not None).
      This parameter is ignored if `initial_population`
      is specified. Used to generate the population from the
      `initial_position` by adding random normal noise with zero mean and
      the specified standard deviation.
    seed: Seed for random number generation.

  Returns:
    A list of `Tensor`s. The initial population.
  """
  if initial_population is not None:
    return [tf.convert_to_tensor(value=part) for part in initial_population]
  # Constructs the population by adding normal noise to the initial position.
  seed_stream = distributions.SeedStream(seed, salt='get_starting_population')
  population = []
  for part in initial_position:
    part = tf.convert_to_tensor(value=part)
    part_event_shape = tf.shape(input=part)
    # We only draw population_size-1 random vectors because we want to ensure
    # that the supplied position is part of the population. The first member
    # is set to be the initial_position.
    population_part_shape = tf.concat([[population_size-1],
                                       part_event_shape], axis=0)
    population_part = tf.random.normal(population_part_shape,
                                       stddev=population_stddev,
                                       dtype=part.dtype.base_dtype,
                                       seed=seed_stream())
    population_part += part
    population_part = tf.concat([[part], population_part], axis=0)
    population.append(population_part)
  return population

Response: Constructs the starting population.


Instruction:
Create a Python 3 function for
performing binary crossover for the current population.

Response:
def _binary_crossover(population,
                      population_size,
                      mutants,
                      crossover_prob,
                      seed):
  """Performs recombination by binary crossover for the current population.

  Let v_i denote the i'th component of the member v and m_i the corresponding
  component of the mutant vector corresponding to v. Then the crossed over
  vector w_i is determined by setting w_i =
  (m_i with probability=crossover_prob else v_i). In addition, DE requires that
  at least one of the components is crossed over (otherwise we end
  up with no change). This is done by choosing on index say k randomly where
  a force crossover is performed (i.e. w_k = m_k). This is the scheme
  implemented in this function.

  Args:
    population: A Python list of `Tensor`s where each `Tensor` in the list
      must be of rank at least 1 and all the elements must have a common
      first dimension. The base population to cross over.
    population_size: A scalar integer `Tensor`. The number of elements in the
      population (i.e. size of the first dimension of any member of
      `population`).
    mutants: A Python list of `Tensor`s with the same structure as `population`.
      The mutated population.
    crossover_prob: A postive real scalar `Tensor` bounded above by 1.0. The
      probability of a crossover being performed for each axis.
    seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
      applied.

  Returns:
    A list of `Tensor`s of the same structure, dtype and shape as `population`.
    The recombined population.
  """
  sizes = [tf.cast(tf.size(input=x), dtype=tf.float64) for x in population]
  seed_stream = distributions.SeedStream(seed, salt='binary_crossover')
  force_crossover_group = distributions.Categorical(sizes).sample(
      [population_size, 1], seed=seed_stream())
  recombinants = []
  for i, population_part in enumerate(population):
    pop_part_flat = tf.reshape(population_part, [population_size, -1])
    mutant_part_flat = tf.reshape(mutants[i], [population_size, -1])
    part_size = tf.size(input=population_part) // population_size
    force_crossovers = tf.one_hot(
        tf.random.uniform([population_size],
                          minval=0,
                          maxval=part_size,
                          dtype=tf.int32,
                          seed=seed_stream()),
        part_size,
        on_value=True,
        off_value=False,
        dtype=tf.bool)  # Tensor of shape [population_size, size]
    group_mask = tf.math.equal(force_crossover_group, i)
    force_crossovers &= group_mask
    do_binary_crossover = tf.random.uniform(
        [population_size, part_size],
        dtype=crossover_prob.dtype.base_dtype,
        seed=seed_stream()) < crossover_prob
    do_binary_crossover |= force_crossovers
    recombinant_flat = tf.where(
        do_binary_crossover,
        x=mutant_part_flat,
        y=pop_part_flat)
    recombinant = tf.reshape(recombinant_flat, tf.shape(input=population_part))
    recombinants.append(recombinant)
  return recombinants


Instruction:
Write a Python 3 script to
compute the mutatated vectors for each population member.

Response:
def _get_mutants(population,
                 population_size,
                 mixing_indices,
                 differential_weight):
  """Computes the mutatated vectors for each population member.

  Args:
    population:  Python `list` of `Tensor`s representing the
      current population vectors. Each `Tensor` must be of the same real dtype.
      The first dimension of each `Tensor` indexes individual
      population members. For example, if the population is a list with a
      single `Tensor` of shape [n, m1, m2], then `n` is the population size and
      the shape of an individual solution is [m1, m2].
      If there is more than one element in the population, then each `Tensor`
      in the list should have the first axis of the same size.
    population_size: Scalar integer `Tensor`. The size of the population.
    mixing_indices: `Tensor` of integral dtype and shape [n, 3] where `n` is the
      number of members in the population. Each element of the `Tensor` must be
      a valid index into the first dimension of the population (i.e range
      between `0` and `n-1` inclusive).
    differential_weight: Real scalar `Tensor`. Must be positive and less than
      2.0. The parameter controlling the strength of mutation.

  Returns:
    mutants: `Tensor` or Python `list` of `Tensor`s of the same shape and dtype
      as the input population. The mutated vectors.
  """
  mixing_indices = tf.reshape(mixing_indices, [-1])
  weights = tf.stack([1.0, differential_weight, -differential_weight])
  def _mutant_part(population_part):
    donors = tf.gather(population_part, mixing_indices)
    donors = tf.transpose(
        a=tf.reshape(donors, [population_size, 3, -1]), perm=[0, 2, 1])
    return tf.math.reduce_sum(input_tensor=donors * weights, axis=-1)

  return [_mutant_part(population_part) for population_part in population]


Instruction:
Explain what the following Python 3 code does
def _get_mixing_indices(size, seed=None, name=None):
  with tf.compat.v1.name_scope(
      name, default_name='get_mixing_indices', values=[size]):
    size = tf.convert_to_tensor(value=size)
    dtype = size.dtype
    seed_stream = distributions.SeedStream(seed, salt='get_mixing_indices')
    first = tf.random.uniform([size],
                              maxval=size-1,
                              dtype=dtype,
                              seed=seed_stream())
    second = tf.random.uniform([size],
                               maxval=size-2,
                               dtype=dtype,
                               seed=seed_stream())
    third = tf.random.uniform([size],
                              maxval=size-3,
                              dtype=dtype,
                              seed=seed_stream())

    # Shift second if it is on top of or to the right of first
    second = tf.where(first < second, x=second, y=second + 1)
    smaller = tf.math.minimum(first, second)
    larger = tf.math.maximum(first, second)
    # Shift the third one so it does not coincide with either the first or the
    # second number. Assuming first < second, shift by 1 if the number is in
    # [first, second) and by 2 if the number is greater than or equal to the
    # second.
    third = tf.where(third < smaller, x=third, y=third + 1)
    third = tf.where(third < larger, x=third, y=third + 1)
    sample = tf.stack([first, second, third], axis=1)
    to_avoid = tf.expand_dims(tf.range(size), axis=-1)
    sample = tf.where(sample < to_avoid, x=sample, y=sample + 1)
    return sample

Response: Generates an array of indices suitable for a differential evolution mutation operation.


Instruction:
Make a summary of the following Python 3 code
def _ensure_list(tensor_or_list):
  if isinstance(tensor_or_list, (list, tuple)):
    return list(tensor_or_list), True
  return [tensor_or_list], False

Response: Converts the input arg to a list if it is not a list already.


Instruction:
Can you generate the documentation for the following Python 3 function
def _get_tol(tol, dtype, validate_args):
  if tol is None:
    return tf.convert_to_tensor(value=0, dtype=dtype)

  tol = tf.convert_to_tensor(value=tol, dtype=dtype)
  if validate_args:
    tol = distribution_util.with_dependencies([
        assert_util.assert_non_negative(
            tol, message="Argument 'tol' must be non-negative")
    ], tol)
  return tol

Response: Gets a Tensor of type dtype 0 if tol is None validation optional.


Instruction:
Write a Python 3 function for
calculating the batched KL divergence with a Deterministic distribution.

Response:
def _kl_deterministic_distribution(a, b, name=None):
  """Calculate the batched KL divergence `KL(a || b)` with `a` Deterministic.

  Args:
    a: instance of a Deterministic distribution object.
    b: instance of a Distribution distribution object.
    name: (optional) Name to use for created operations. Default is
      "kl_deterministic_distribution".

  Returns:
    Batchwise `KL(a || b)`.
  """
  with tf.name_scope(name or "kl_deterministic_distribution"):
    return -b.log_prob(a.loc)


Instruction:
Implement a Python 3 function for
clipping values to a specified min and max while leaving gradient unaltered.

Response:
def clip_by_value_preserve_gradient(t, clip_value_min, clip_value_max,
                                    name=None):
  """Clips values to a specified min and max while leaving gradient unaltered.

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

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

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

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

  Returns:
    clipped_t: A clipped `Tensor`.
  """
  with tf.compat.v1.name_scope(name, 'clip_by_value_preserve_gradient',
                               [t, clip_value_min, clip_value_max]):
    t = tf.convert_to_tensor(value=t, name='t')
    clip_t = tf.clip_by_value(t, clip_value_min, clip_value_max)
    return t + tf.stop_gradient(clip_t - t)


Instruction:
Can you generate the documentation for the following Python 3 function
def build_input_pipeline(train_images, batch_size):

  training_dataset = tf.data.Dataset.from_tensor_slices(train_images)
  training_batches = training_dataset.shuffle(
      50000, reshuffle_each_iteration=True).repeat().batch(batch_size)
  training_iterator = tf.compat.v1.data.make_one_shot_iterator(training_batches)
  images = training_iterator.get_next()
  return images

Response: Build an iterator over training batches.


Instruction:
Write a Python 3 function that can
save a synthetic image as a PNG file.

Response:
def plot_generated_images(images, fname):
  """Save a synthetic image as a PNG file.

  Args:
    images: samples of synthetic images generated by the generative network.
    fname: Python `str`, filename to save the plot to.
  """
  fig = plt.figure(figsize=(4, 4))
  canvas = backend_agg.FigureCanvasAgg(fig)

  for i, image in enumerate(images):
    ax = fig.add_subplot(4, 4, i + 1)
    plt.axis('off')
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    ax.imshow(image.reshape(IMAGE_SHAPE[:-1]), cmap='Greys_r')

  fig.tight_layout()
  plt.subplots_adjust(wspace=0.05, hspace=0.05)
  canvas.print_figure(fname, format='png')


Instruction:
Create a Python 3 function for
converting a sequence of productions into a string of terminal symbols.

Response:
def convert_to_string(self, productions):
    """Converts a sequence of productions into a string of terminal symbols.

    Args:
      productions: Tensor of shape [1, num_productions, num_production_rules].
        Slices along the `num_productions` dimension represent one-hot vectors.

    Returns:
      str that concatenates all terminal symbols from `productions`.

    Raises:
      ValueError: If the first production rule does not begin with
        `self.start_symbol`.
    """
    symbols = []
    for production in tf.unstack(productions, axis=1):
      lhs, rhs = self.production_rules[tf.argmax(input=production, axis=-1)]
      if not symbols:  # first iteration
        if lhs != self.start_symbol:
          raise ValueError("`productions` must begin with `self.start_symbol`.")
        symbols = rhs
      else:
        # Greedily unroll the nonterminal symbols based on the first occurrence
        # in a linear sequence.
        index = symbols.index(lhs)
        symbols = symbols[:index] + rhs + symbols[index + 1:]
    string = "".join(symbols)
    return string


Instruction:
How would you code a function in Python 3 to
produce a masking tensor for valid production rules.

Response:
def mask(self, symbol, on_value, off_value):
    """Produces a masking tensor for (in)valid production rules.

    Args:
      symbol: str, a symbol in the grammar.
      on_value: Value to use for a valid production rule.
      off_value: Value to use for an invalid production rule.

    Returns:
      Tensor of shape [1, num_production_rules]. An element is `on_value`
      if its corresponding production rule has `symbol` on its left-hand-side;
      the element is `off_value` otherwise.
    """
    mask_values = [on_value if lhs == symbol else off_value
                   for lhs, _ in self.production_rules]
    mask_values = tf.reshape(mask_values, [1, len(self.production_rules)])
    return mask_values


Instruction:
Make a summary of the following Python 3 code
def call(self, inputs):
    del inputs  # unused
    latent_code = ed.MultivariateNormalDiag(loc=tf.zeros(self.latent_size),
                                            sample_shape=1,
                                            name="latent_code")
    state = self.lstm.zero_state(1, dtype=tf.float32)
    t = 0
    productions = []
    stack = [self.grammar.start_symbol]
    while stack:
      symbol = stack.pop()
      net, state = self.lstm(latent_code, state)
      logits = (self.output_layer(net) +
                self.grammar.mask(symbol, on_value=0., off_value=-1e9))
      production = ed.OneHotCategorical(logits=logits,
                                        name="production_" + str(t))
      _, rhs = self.grammar.production_rules[tf.argmax(
          input=production, axis=-1)]
      for symbol in rhs:
        if symbol in self.grammar.nonterminal_symbols:
          stack.append(symbol)
      productions.append(production)
      t += 1
    return tf.stack(productions, axis=1)

Response: Runs the model forward to generate a sequence of productions.


Instruction:
Explain what the following Python 3 code does
def call(self, inputs):
    net = self.encoder_net(tf.cast(inputs, tf.float32))
    return ed.MultivariateNormalDiag(
        loc=net[..., :self.latent_size],
        scale_diag=tf.nn.softplus(net[..., self.latent_size:]),
        name="latent_code_posterior")

Response: Runs the model forward to return a stochastic encoding.


Instruction:
Explain what the following Python 3 code does
def _hat_integral(self, x):
    x = tf.cast(x, self.power.dtype)
    t = self.power - 1.
    return tf.exp((-t) * tf.math.log1p(x) - tf.math.log(t))

Response: Integral of the hat function used for sampling.


Instruction:
Here you have a function in Python 3, explain what it does
def _hat_integral_inverse(self, x):
    x = tf.cast(x, self.power.dtype)
    t = self.power - 1.
    return tf.math.expm1(-(tf.math.log(t) + tf.math.log(x)) / t)

Response: Inverse function of _hat_integral.


Instruction:
Given the following Python 3 function, write the documentation
def matrix_rank(a, tol=None, validate_args=False, name=None):
  with tf.compat.v1.name_scope(name, 'matrix_rank', [a, tol]):
    a = tf.convert_to_tensor(value=a, dtype_hint=tf.float32, name='a')
    assertions = _maybe_validate_matrix(a, validate_args)
    if assertions:
      with tf.control_dependencies(assertions):
        a = tf.identity(a)
    s = tf.linalg.svd(a, compute_uv=False)
    if tol is None:
      if a.shape[-2:].is_fully_defined():
        m = np.max(a.shape[-2:].as_list())
      else:
        m = tf.reduce_max(input_tensor=tf.shape(input=a)[-2:])
      eps = np.finfo(a.dtype.as_numpy_dtype).eps
      tol = (eps * tf.cast(m, a.dtype) *
             tf.reduce_max(input_tensor=s, axis=-1, keepdims=True))
    return tf.reduce_sum(input_tensor=tf.cast(s > tol, tf.int32), axis=-1)

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


Instruction:
How would you code a function in Python 3 to
concatenate cholesky decomposition of chol.

Response:
def cholesky_concat(chol, cols, name=None):
  """Concatenates `chol @ chol.T` with additional rows and columns.

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

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

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

  Returns:
    chol_concat: The Cholesky decomposition of:
      ```
      [ [ mat  cols[:n, :] ]
        [   conj(cols.T)   ] ]
      ```
  """
  with tf.compat.v2.name_scope(name or 'cholesky_extend'):
    dtype = dtype_util.common_dtype([chol, cols], preferred_dtype=tf.float32)
    chol = tf.convert_to_tensor(value=chol, name='chol', dtype=dtype)
    cols = tf.convert_to_tensor(value=cols, name='cols', dtype=dtype)
    n = prefer_static.shape(chol)[-1]
    mat_nm, mat_mm = cols[..., :n, :], cols[..., n:, :]
    solved_nm = linear_operator_util.matrix_triangular_solve_with_broadcast(
        chol, mat_nm)
    lower_right_mm = tf.linalg.cholesky(
        mat_mm - tf.matmul(solved_nm, solved_nm, adjoint_a=True))
    lower_left_mn = tf.math.conj(tf.linalg.matrix_transpose(solved_nm))
    out_batch = prefer_static.shape(solved_nm)[:-2]
    chol = tf.broadcast_to(
        chol,
        tf.concat([out_batch, prefer_static.shape(chol)[-2:]], axis=0))
    top_right_zeros_nm = tf.zeros_like(solved_nm)
    return tf.concat([tf.concat([chol, top_right_zeros_nm], axis=-1),
                      tf.concat([lower_left_mn, lower_right_mm], axis=-1)],
                     axis=-2)


Instruction:
Write a Python 3 script for
swapping `m` and `i` on axis -1. (Helper for pivoted_cholesky.) Given a batch of int64 vectors `vecs`, scalar index `m`, and compatibly shaped per-vector indices `i`, this function swaps elements `m` and `i` in each vector. For the use-case below, these are permutation vectors. Args: vecs: Vectors on which we perform the swap, int64 `Tensor`. m: Scalar int64 `Tensor`, the index into which the `i`th element is going. i: Batch int64 `Tensor`, shaped like vecs.shape[:-1] + [1]; the index into which the `m`th element is going. Returns: vecs: The updated vectors.

Response:
def _swap_m_with_i(vecs, m, i):
  """Swaps `m` and `i` on axis -1. (Helper for pivoted_cholesky.)

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

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

  Returns:
    vecs: The updated vectors.
  """
  vecs = tf.convert_to_tensor(value=vecs, dtype=tf.int64, name='vecs')
  m = tf.convert_to_tensor(value=m, dtype=tf.int64, name='m')
  i = tf.convert_to_tensor(value=i, dtype=tf.int64, name='i')
  trailing_elts = tf.broadcast_to(
      tf.range(m + 1,
               prefer_static.shape(vecs, out_type=tf.int64)[-1]),
      prefer_static.shape(vecs[..., m + 1:]))
  shp = prefer_static.shape(trailing_elts)
  trailing_elts = tf.where(
      tf.equal(trailing_elts, tf.broadcast_to(i, shp)),
      tf.broadcast_to(tf.gather(vecs, [m], axis=-1), shp),
      tf.broadcast_to(vecs[..., m + 1:], shp))
  # TODO(bjp): Could we use tensor_scatter_nd_update?
  vecs_shape = vecs.shape
  vecs = tf.concat([
      vecs[..., :m],
      tf.gather(vecs, i, batch_dims=prefer_static.rank(vecs) - 1), trailing_elts
  ], axis=-1)
  tensorshape_util.set_shape(vecs, vecs_shape)
  return vecs


Instruction:
Can you generate the documentation for the following Python 3 function
def pivoted_cholesky(matrix, max_rank, diag_rtol=1e-3, name=None):
  with tf.compat.v2.name_scope(name or 'pivoted_cholesky'):
    dtype = dtype_util.common_dtype([matrix, diag_rtol],
                                    preferred_dtype=tf.float32)
    matrix = tf.convert_to_tensor(value=matrix, name='matrix', dtype=dtype)
    if tensorshape_util.rank(matrix.shape) is None:
      raise NotImplementedError('Rank of `matrix` must be known statically')

    max_rank = tf.convert_to_tensor(
        value=max_rank, name='max_rank', dtype=tf.int64)
    max_rank = tf.minimum(max_rank,
                          prefer_static.shape(matrix, out_type=tf.int64)[-1])
    diag_rtol = tf.convert_to_tensor(
        value=diag_rtol, dtype=dtype, name='diag_rtol')
    matrix_diag = tf.linalg.diag_part(matrix)
    # matrix is P.D., therefore all matrix_diag > 0, so we don't need abs.
    orig_error = tf.reduce_max(input_tensor=matrix_diag, axis=-1)

    def cond(m, pchol, perm, matrix_diag):
      """Condition for `tf.while_loop` continuation."""
      del pchol
      del perm
      error = tf.linalg.norm(tensor=matrix_diag, ord=1, axis=-1)
      max_err = tf.reduce_max(input_tensor=error / orig_error)
      return (m < max_rank) & (tf.equal(m, 0) | (max_err > diag_rtol))

    batch_dims = tensorshape_util.rank(matrix.shape) - 2
    def batch_gather(params, indices, axis=-1):
      return tf.gather(params, indices, axis=axis, batch_dims=batch_dims)

    def body(m, pchol, perm, matrix_diag):
      """Body of a single `tf.while_loop` iteration."""
      # Here is roughly a numpy, non-batched version of what's going to happen.
      # (See also Algorithm 1 of Harbrecht et al.)
      # 1: maxi = np.argmax(matrix_diag[perm[m:]]) + m
      # 2: maxval = matrix_diag[perm][maxi]
      # 3: perm[m], perm[maxi] = perm[maxi], perm[m]
      # 4: row = matrix[perm[m]][perm[m + 1:]]
      # 5: row -= np.sum(pchol[:m][perm[m + 1:]] * pchol[:m][perm[m]]], axis=-2)
      # 6: pivot = np.sqrt(maxval); row /= pivot
      # 7: row = np.concatenate([[[pivot]], row], -1)
      # 8: matrix_diag[perm[m:]] -= row**2
      # 9: pchol[m, perm[m:]] = row

      # Find the maximal position of the (remaining) permuted diagonal.
      # Steps 1, 2 above.
      permuted_diag = batch_gather(matrix_diag, perm[..., m:])
      maxi = tf.argmax(
          input=permuted_diag, axis=-1, output_type=tf.int64)[..., tf.newaxis]
      maxval = batch_gather(permuted_diag, maxi)
      maxi = maxi + m
      maxval = maxval[..., 0]
      # Update perm: Swap perm[...,m] with perm[...,maxi]. Step 3 above.
      perm = _swap_m_with_i(perm, m, maxi)
      # Step 4.
      row = batch_gather(matrix, perm[..., m:m + 1], axis=-2)
      row = batch_gather(row, perm[..., m + 1:])
      # Step 5.
      prev_rows = pchol[..., :m, :]
      prev_rows_perm_m_onward = batch_gather(prev_rows, perm[..., m + 1:])
      prev_rows_pivot_col = batch_gather(prev_rows, perm[..., m:m + 1])
      row -= tf.reduce_sum(
          input_tensor=prev_rows_perm_m_onward * prev_rows_pivot_col,
          axis=-2)[..., tf.newaxis, :]
      # Step 6.
      pivot = tf.sqrt(maxval)[..., tf.newaxis, tf.newaxis]
      # Step 7.
      row = tf.concat([pivot, row / pivot], axis=-1)
      # TODO(b/130899118): Pad grad fails with int64 paddings.
      # Step 8.
      paddings = tf.concat([
          tf.zeros([prefer_static.rank(pchol) - 1, 2], dtype=tf.int32),
          [[tf.cast(m, tf.int32), 0]]], axis=0)
      diag_update = tf.pad(tensor=row**2, paddings=paddings)[..., 0, :]
      reverse_perm = _invert_permutation(perm)
      matrix_diag -= batch_gather(diag_update, reverse_perm)
      # Step 9.
      row = tf.pad(tensor=row, paddings=paddings)
      # TODO(bjp): Defer the reverse permutation all-at-once at the end?
      row = batch_gather(row, reverse_perm)
      pchol_shape = pchol.shape
      pchol = tf.concat([pchol[..., :m, :], row, pchol[..., m + 1:, :]],
                        axis=-2)
      tensorshape_util.set_shape(pchol, pchol_shape)
      return m + 1, pchol, perm, matrix_diag

    m = np.int64(0)
    pchol = tf.zeros_like(matrix[..., :max_rank, :])
    matrix_shape = prefer_static.shape(matrix, out_type=tf.int64)
    perm = tf.broadcast_to(
        prefer_static.range(matrix_shape[-1]), matrix_shape[:-1])
    _, pchol, _, _ = tf.while_loop(
        cond=cond, body=body, loop_vars=(m, pchol, perm, matrix_diag))
    pchol = tf.linalg.matrix_transpose(pchol)
    tensorshape_util.set_shape(
        pchol, tensorshape_util.concatenate(matrix_diag.shape, [None]))
    return pchol

Response: Computes the ( partial ) pivoted Cholesky decomposition of matrix.


Instruction:
Given the following Python 3 function, write the documentation
def pinv(a, rcond=None, validate_args=False, name=None):
  with tf.compat.v1.name_scope(name, 'pinv', [a, rcond]):
    a = tf.convert_to_tensor(value=a, name='a')

    assertions = _maybe_validate_matrix(a, validate_args)
    if assertions:
      with tf.control_dependencies(assertions):
        a = tf.identity(a)

    dtype = a.dtype.as_numpy_dtype

    if rcond is None:
      def get_dim_size(dim):
        if tf.compat.dimension_value(a.shape[dim]) is not None:
          return tf.compat.dimension_value(a.shape[dim])
        return tf.shape(input=a)[dim]

      num_rows = get_dim_size(-2)
      num_cols = get_dim_size(-1)
      if isinstance(num_rows, int) and isinstance(num_cols, int):
        max_rows_cols = float(max(num_rows, num_cols))
      else:
        max_rows_cols = tf.cast(tf.maximum(num_rows, num_cols), dtype)
      rcond = 10. * max_rows_cols * np.finfo(dtype).eps

    rcond = tf.convert_to_tensor(value=rcond, dtype=dtype, name='rcond')

    # Calculate pseudo inverse via SVD.
    # Note: if a is symmetric then u == v. (We might observe additional
    # performance by explicitly setting `v = u` in such cases.)
    [
        singular_values,         # Sigma
        left_singular_vectors,   # U
        right_singular_vectors,  # V
    ] = tf.linalg.svd(a, full_matrices=False, compute_uv=True)

    # Saturate small singular values to inf. This has the effect of make
    # `1. / s = 0.` while not resulting in `NaN` gradients.
    cutoff = rcond * tf.reduce_max(input_tensor=singular_values, axis=-1)
    singular_values = tf.where(
        singular_values > cutoff[..., tf.newaxis], singular_values,
        tf.fill(tf.shape(input=singular_values), np.array(np.inf, dtype)))

    # Although `a == tf.matmul(u, s * v, transpose_b=True)` we swap
    # `u` and `v` here so that `tf.matmul(pinv(A), A) = tf.eye()`, i.e.,
    # a matrix inverse has "transposed" semantics.
    a_pinv = tf.matmul(
        right_singular_vectors / singular_values[..., tf.newaxis, :],
        left_singular_vectors,
        adjoint_b=True)

    if a.shape.ndims is not None:
      a_pinv.set_shape(a.shape[:-2].concatenate([a.shape[-1], a.shape[-2]]))

    return a_pinv

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


Instruction:
Can you tell what is the following Python 3 function doing
def lu_solve(lower_upper, perm, rhs,
             validate_args=False,
             name=None):
  """Solves systems of linear eqns `A X = RHS`, given LU factorizations.

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

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

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

  #### Examples

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

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

  """

  with tf.compat.v1.name_scope(name, 'lu_solve', [lower_upper, perm, rhs]):
    lower_upper = tf.convert_to_tensor(
        value=lower_upper, dtype_hint=tf.float32, name='lower_upper')
    perm = tf.convert_to_tensor(value=perm, dtype_hint=tf.int32, name='perm')
    rhs = tf.convert_to_tensor(
        value=rhs, dtype_hint=lower_upper.dtype, name='rhs')

    assertions = _lu_solve_assertions(lower_upper, perm, rhs, validate_args)
    if assertions:
      with tf.control_dependencies(assertions):
        lower_upper = tf.identity(lower_upper)
        perm = tf.identity(perm)
        rhs = tf.identity(rhs)

    if rhs.shape.ndims == 2 and perm.shape.ndims == 1:
      # Both rhs and perm have scalar batch_shape.
      permuted_rhs = tf.gather(rhs, perm, axis=-2)
    else:
      # Either rhs or perm have non-scalar batch_shape or we can't determine
      # this information statically.
      rhs_shape = tf.shape(input=rhs)
      broadcast_batch_shape = tf.broadcast_dynamic_shape(
          rhs_shape[:-2],
          tf.shape(input=perm)[:-1])
      d, m = rhs_shape[-2], rhs_shape[-1]
      rhs_broadcast_shape = tf.concat([broadcast_batch_shape, [d, m]], axis=0)

      # Tile out rhs.
      broadcast_rhs = tf.broadcast_to(rhs, rhs_broadcast_shape)
      broadcast_rhs = tf.reshape(broadcast_rhs, [-1, d, m])

      # Tile out perm and add batch indices.
      broadcast_perm = tf.broadcast_to(perm, rhs_broadcast_shape[:-1])
      broadcast_perm = tf.reshape(broadcast_perm, [-1, d])
      broadcast_batch_size = tf.reduce_prod(input_tensor=broadcast_batch_shape)
      broadcast_batch_indices = tf.broadcast_to(
          tf.range(broadcast_batch_size)[:, tf.newaxis],
          [broadcast_batch_size, d])
      broadcast_perm = tf.stack([broadcast_batch_indices, broadcast_perm],
                                axis=-1)

      permuted_rhs = tf.gather_nd(broadcast_rhs, broadcast_perm)
      permuted_rhs = tf.reshape(permuted_rhs, rhs_broadcast_shape)

    lower = tf.linalg.set_diag(
        tf.linalg.band_part(lower_upper, num_lower=-1, num_upper=0),
        tf.ones(tf.shape(input=lower_upper)[:-1], dtype=lower_upper.dtype))
    return linear_operator_util.matrix_triangular_solve_with_broadcast(
        lower_upper,  # Only upper is accessed.
        linear_operator_util.matrix_triangular_solve_with_broadcast(
            lower, permuted_rhs),
        lower=False)

Response: Solves system of linear eqns A X = RHS given LU factorizations.


Instruction:
Create a Python 3 function to
compute the matrix inverse given the matrix s LU decomposition.

Response:
def lu_matrix_inverse(lower_upper, perm, validate_args=False, name=None):
  """Computes a matrix inverse given the matrix's LU decomposition.

  This op is conceptually identical to,

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

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

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

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

  #### Examples

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

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

  """

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


Instruction:
Make a summary of the following Python 3 code
def lu_reconstruct(lower_upper, perm, validate_args=False, name=None):
  with tf.compat.v1.name_scope(name, 'lu_reconstruct', [lower_upper, perm]):
    lower_upper = tf.convert_to_tensor(
        value=lower_upper, dtype_hint=tf.float32, name='lower_upper')
    perm = tf.convert_to_tensor(value=perm, dtype_hint=tf.int32, name='perm')

    assertions = _lu_reconstruct_assertions(lower_upper, perm, validate_args)
    if assertions:
      with tf.control_dependencies(assertions):
        lower_upper = tf.identity(lower_upper)
        perm = tf.identity(perm)

    shape = tf.shape(input=lower_upper)

    lower = tf.linalg.set_diag(
        tf.linalg.band_part(lower_upper, num_lower=-1, num_upper=0),
        tf.ones(shape[:-1], dtype=lower_upper.dtype))
    upper = tf.linalg.band_part(lower_upper, num_lower=0, num_upper=-1)
    x = tf.matmul(lower, upper)

    if lower_upper.shape.ndims is None or lower_upper.shape.ndims != 2:
      # We either don't know the batch rank or there are >0 batch dims.
      batch_size = tf.reduce_prod(input_tensor=shape[:-2])
      d = shape[-1]
      x = tf.reshape(x, [batch_size, d, d])
      perm = tf.reshape(perm, [batch_size, d])
      perm = tf.map_fn(tf.math.invert_permutation, perm)
      batch_indices = tf.broadcast_to(
          tf.range(batch_size)[:, tf.newaxis],
          [batch_size, d])
      x = tf.gather_nd(x, tf.stack([batch_indices, perm], axis=-1))
      x = tf.reshape(x, shape)
    else:
      x = tf.gather(x, tf.math.invert_permutation(perm))

    x.set_shape(lower_upper.shape)
    return x

Response: The inverse LU decomposition of X.


Instruction:
Can you write a function in Python 3 where it
returns a list of assertions related to lu_reconstruct assumptions.

Response:
def _lu_reconstruct_assertions(lower_upper, perm, validate_args):
  """Returns list of assertions related to `lu_reconstruct` assumptions."""
  assertions = []

  message = 'Input `lower_upper` must have at least 2 dimensions.'
  if lower_upper.shape.ndims is not None:
    if lower_upper.shape.ndims < 2:
      raise ValueError(message)
  elif validate_args:
    assertions.append(
        tf.compat.v1.assert_rank_at_least(lower_upper, rank=2, message=message))

  message = '`rank(lower_upper)` must equal `rank(perm) + 1`'
  if lower_upper.shape.ndims is not None and perm.shape.ndims is not None:
    if lower_upper.shape.ndims != perm.shape.ndims + 1:
      raise ValueError(message)
  elif validate_args:
    assertions.append(
        tf.compat.v1.assert_rank(
            lower_upper, rank=tf.rank(perm) + 1, message=message))

  message = '`lower_upper` must be square.'
  if lower_upper.shape[:-2].is_fully_defined():
    if lower_upper.shape[-2] != lower_upper.shape[-1]:
      raise ValueError(message)
  elif validate_args:
    m, n = tf.split(tf.shape(input=lower_upper)[-2:], num_or_size_splits=2)
    assertions.append(tf.compat.v1.assert_equal(m, n, message=message))

  return assertions


Instruction:
Here you have a function in Python 3, explain what it does
def _lu_solve_assertions(lower_upper, perm, rhs, validate_args):
  assertions = _lu_reconstruct_assertions(lower_upper, perm, validate_args)

  message = 'Input `rhs` must have at least 2 dimensions.'
  if rhs.shape.ndims is not None:
    if rhs.shape.ndims < 2:
      raise ValueError(message)
  elif validate_args:
    assertions.append(
        tf.compat.v1.assert_rank_at_least(rhs, rank=2, message=message))

  message = '`lower_upper.shape[-1]` must equal `rhs.shape[-1]`.'
  if (tf.compat.dimension_value(lower_upper.shape[-1]) is not None and
      tf.compat.dimension_value(rhs.shape[-2]) is not None):
    if lower_upper.shape[-1] != rhs.shape[-2]:
      raise ValueError(message)
  elif validate_args:
    assertions.append(
        tf.compat.v1.assert_equal(
            tf.shape(input=lower_upper)[-1],
            tf.shape(input=rhs)[-2],
            message=message))

  return assertions

Response: Returns list of assertions related to lu_solve assumptions.


Instruction:
Write a Python 3 script to
return a Tensor or SparseTensor with a Tensor.

Response:
def sparse_or_dense_matmul(sparse_or_dense_a,
                           dense_b,
                           validate_args=False,
                           name=None,
                           **kwargs):
  """Returns (batched) matmul of a SparseTensor (or Tensor) with a Tensor.

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

  Returns:
    product: A dense (batch of) matrix-shaped Tensor of the same batch shape and
    dtype as `sparse_or_dense_a` and `dense_b`. If `sparse_or_dense_a` or
    `dense_b` is adjointed through `kwargs` then the shape is adjusted
    accordingly.
  """
  with tf.compat.v1.name_scope(name, 'sparse_or_dense_matmul',
                               [sparse_or_dense_a, dense_b]):
    dense_b = tf.convert_to_tensor(
        value=dense_b, dtype_hint=tf.float32, name='dense_b')

    if validate_args:
      assert_a_rank_at_least_2 = tf.compat.v1.assert_rank_at_least(
          sparse_or_dense_a,
          rank=2,
          message='Input `sparse_or_dense_a` must have at least 2 dimensions.')
      assert_b_rank_at_least_2 = tf.compat.v1.assert_rank_at_least(
          dense_b,
          rank=2,
          message='Input `dense_b` must have at least 2 dimensions.')
      with tf.control_dependencies(
          [assert_a_rank_at_least_2, assert_b_rank_at_least_2]):
        sparse_or_dense_a = tf.identity(sparse_or_dense_a)
        dense_b = tf.identity(dense_b)

    if isinstance(sparse_or_dense_a,
                  (tf.SparseTensor, tf.compat.v1.SparseTensorValue)):
      return _sparse_tensor_dense_matmul(sparse_or_dense_a, dense_b, **kwargs)
    else:
      return tf.matmul(sparse_or_dense_a, dense_b, **kwargs)


Instruction:
How would you implement a function in Python 3 that
returns a Tensor of the same shape as sparse_or_dense_matrix and dense_vector.

Response:
def sparse_or_dense_matvecmul(sparse_or_dense_matrix,
                              dense_vector,
                              validate_args=False,
                              name=None,
                              **kwargs):
  """Returns (batched) matmul of a (sparse) matrix with a column vector.

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

  Returns:
    product: A dense (batch of) vector-shaped Tensor of the same batch shape and
    dtype as `sparse_or_dense_matrix` and `dense_vector`.
  """
  with tf.compat.v1.name_scope(name, 'sparse_or_dense_matvecmul',
                               [sparse_or_dense_matrix, dense_vector]):
    dense_vector = tf.convert_to_tensor(
        value=dense_vector, dtype_hint=tf.float32, name='dense_vector')
    return tf.squeeze(
        sparse_or_dense_matmul(
            sparse_or_dense_matrix,
            dense_vector[..., tf.newaxis],
            validate_args=validate_args,
            **kwargs),
        axis=[-1])


Instruction:
Write a Python 3 function that can
return a dense matrix - shaped Tensor with a Tensor.

Response:
def _sparse_tensor_dense_matmul(sp_a, b, **kwargs):
  """Returns (batched) matmul of a SparseTensor with a Tensor.

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

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

  # Reshape the SparseTensor into a rank 3 SparseTensors, with the
  # batch shape flattened to a single dimension. If the batch rank is 0, then
  # we add a batch dimension of rank 1.
  sp_a = tf.sparse.reshape(sp_a, tf.concat([[-1], _get_shape(sp_a)[-2:]],
                                           axis=0))
  # Reshape b to stack the batch dimension along the rows.
  b = tf.reshape(b, tf.concat([[-1], _get_shape(b)[-1:]], axis=0))

  # Convert the SparseTensor to a matrix in block diagonal form with blocks of
  # matrices [M, N]. This allow us to use tf.sparse_tensor_dense_matmul which
  # only accepts rank 2 (Sparse)Tensors.
  out = tf.sparse.sparse_dense_matmul(_sparse_block_diag(sp_a), b, **kwargs)

  # Finally retrieve the original batch shape from the resulting rank 2 Tensor.
  # Note that we avoid inferring the final shape from `sp_a` or `b` because we
  # might have transposed one or both of them.
  return tf.reshape(
      out,
      tf.concat([batch_shape, [-1], _get_shape(out)[-1:]], axis=0))


Instruction:
Implement a function in Python 3 to
return a SparseTensor from a batch of SparseTensors.

Response:
def _sparse_block_diag(sp_a):
  """Returns a block diagonal rank 2 SparseTensor from a batch of SparseTensors.

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

  Returns:
    sp_block_diag_a: matrix-shaped, `float` `SparseTensor` with the same dtype
    as `sparse_or_matrix`, of shape [B * M, B * N] where `sp_a` has shape
    [B, M, N]. Each [M, N] batch of `sp_a` is lined up along the diagonal.
  """
  # Construct the matrix [[M, N], [1, 0], [0, 1]] which would map the index
  # (b, i, j) to (Mb + i, Nb + j). This effectively creates a block-diagonal
  # matrix of dense shape [B * M, B * N].
  # Note that this transformation doesn't increase the number of non-zero
  # entries in the SparseTensor.
  sp_a_shape = tf.convert_to_tensor(value=_get_shape(sp_a, tf.int64))
  ind_mat = tf.concat([[sp_a_shape[-2:]], tf.eye(2, dtype=tf.int64)], axis=0)
  indices = tf.matmul(sp_a.indices, ind_mat)
  dense_shape = sp_a_shape[0] * sp_a_shape[1:]
  return tf.SparseTensor(
      indices=indices, values=sp_a.values, dense_shape=dense_shape)


Instruction:
Can you tell what is the following Python 3 function doing
def _maybe_validate_matrix(a, validate_args):
  assertions = []
  if not a.dtype.is_floating:
    raise TypeError('Input `a` must have `float`-like `dtype` '
                    '(saw {}).'.format(a.dtype.name))
  if a.shape.ndims is not None:
    if a.shape.ndims < 2:
      raise ValueError('Input `a` must have at least 2 dimensions '
                       '(saw: {}).'.format(a.shape.ndims))
  elif validate_args:
    assertions.append(tf.compat.v1.assert_rank_at_least(
        a, rank=2, message='Input `a` must have at least 2 dimensions.'))
  return assertions

Response: Checks that input is a float matrix.


Instruction:
Can you generate the documentation for the following Python 3 function
def _grad_neg_log_likelihood_and_fim(model_matrix, linear_response, response,
                                     model):
  """Computes the neg-log-likelihood gradient and Fisher information for a GLM.

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

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

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

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

  Returns:
    grad_neg_log_likelihood: (Batch of) vector-shaped `Tensor` with the same
      shape and dtype as a single row of `model_matrix`, representing the
      gradient of the negative log likelihood of `response` given linear
      response `linear_response`.
    fim_middle: (Batch of) vector-shaped `Tensor` with the same shape and dtype
      as a single column of `model_matrix`, satisfying the equation
      `Fisher information =
      Transpose(model_matrix)
      @ diag(fim_middle)
      @ model_matrix`.
  """
  # TODO(b/111926503): Determine whether there are some practical cases where it
  # is computationally favorable to compute the full FIM.
  mean, variance, grad_mean = model(linear_response)

  is_valid = (
      tf.math.is_finite(grad_mean) & tf.not_equal(grad_mean, 0.)
      & tf.math.is_finite(variance) & (variance > 0.))

  def _mask_if_invalid(x, mask):
    mask = tf.fill(
        tf.shape(input=x), value=np.array(mask, x.dtype.as_numpy_dtype))
    return tf.where(is_valid, x, mask)

  # TODO(b/111923449): Link to derivation once it's available.
  v = (response - mean) * _mask_if_invalid(grad_mean, 1) / _mask_if_invalid(
      variance, np.inf)
  grad_log_likelihood = sparse_or_dense_matvecmul(
      model_matrix, v, adjoint_a=True)
  fim_middle = _mask_if_invalid(grad_mean, 0.)**2 / _mask_if_invalid(
      variance, np.inf)
  return -grad_log_likelihood, fim_middle

Response: Returns the gradient of the neg - log - likelihood of a generalized linear model and the Fisher information for a GLM.


Instruction:
Can you create a Python 3 function that
generates the slices for building an autoregressive mask.

Response:
def _gen_slices(num_blocks, n_in, n_out, mask_type=MASK_EXCLUSIVE):
  """Generate the slices for building an autoregressive mask."""
  # TODO(b/67594795): Better support of dynamic shape.
  slices = []
  col = 0
  d_in = n_in // num_blocks
  d_out = n_out // num_blocks
  row = d_out if mask_type == MASK_EXCLUSIVE else 0
  for _ in range(num_blocks):
    row_slice = slice(row, None)
    col_slice = slice(col, col + d_in)
    slices.append([row_slice, col_slice])
    col += d_in
    row += d_out
  return slices


Instruction:
Given the following Python 3 function, write the documentation
def _gen_mask(num_blocks,
              n_in,
              n_out,
              mask_type=MASK_EXCLUSIVE,
              dtype=tf.float32):
  """Generate the mask for building an autoregressive dense layer."""
  # TODO(b/67594795): Better support of dynamic shape.
  mask = np.zeros([n_out, n_in], dtype=dtype.as_numpy_dtype())
  slices = _gen_slices(num_blocks, n_in, n_out, mask_type=mask_type)
  for [row_slice, col_slice] in slices:
    mask[row_slice, col_slice] = 1
  return mask

Response: Generate the mask for building an autoregressive dense layer.


Instruction:
Implement a Python 3 function for
building the Masked Autoregressive Density Estimator (Germain et al., 2015). This will be wrapped in a make_template to ensure the variables are only created once. It takes the input and returns the `loc` ("mu" in [Germain et al. (2015)][1]) and `log_scale` ("alpha" in [Germain et al. (2015)][1]) from the MADE network. Warning: This function uses `masked_dense` to create randomly initialized `tf.Variables`. It is presumed that these will be fit, just as you would any other neural architecture which uses `tf.layers.dense`. #### About Hidden Layers Each element of `hidden_layers` should be greater than the `input_depth` (i.e., `input_depth = tf.shape(input)[-1]` where `input` is the input to the neural network). This is necessary to ensure the autoregressivity property. #### About Clipping This function also optionally clips the `log_scale` (but possibly not its gradient). This is useful because if `log_scale` is too small/large it might underflow/overflow making it impossible for the `MaskedAutoregressiveFlow` bijector to implement a bijection. Additionally, the `log_scale_clip_gradient` `bool` indicates whether the gradient should also be clipped. The default does not clip the gradient; this is useful because it still provides gradient information (for fitting) yet solves the numerical stability problem. I.e., `log_scale_clip_gradient = False` means `grad[exp(clip(x))] = grad[x] exp(clip(x))` rather than the usual `grad[clip(x)] exp(clip(x))`. Args: 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. Default: `False`. activation: Activation function (callable). Explicitly setting to `None` implies a linear activation. log_scale_min_clip: `float`-like scalar `Tensor`, or a `Tensor` with the same shape as `log_scale`. The minimum value to clip by. Default: -5. log_scale_max_clip: `float`-like scalar `Tensor`, or a `Tensor` with the same shape as `log_scale`. The maximum value to clip by. Default: 3. log_scale_clip_gradient: Python `bool` indicating that the gradient of `tf.clip_by_value` should be preserved. Default: `False`. name: A name for ops managed by this function. Default: "masked_autoregressive_default_template". *args: `tf.layers.dense` arguments. **kwargs: `tf.layers.dense` keyword arguments. Returns: shift: `Float`-like `Tensor` of shift terms (the "mu" in [Germain et al. (2015)][1]). log_scale: `Float`-like `Tensor` of log(scale) terms (the "alpha" in [Germain et al. (2015)][1]). Raises: NotImplementedError: if rightmost dimension of `inputs` is unknown prior to graph execution. #### References [1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE: Masked Autoencoder for Distribution Estimation. In _International Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509

Response:
def masked_autoregressive_default_template(hidden_layers,
                                           shift_only=False,
                                           activation=tf.nn.relu,
                                           log_scale_min_clip=-5.,
                                           log_scale_max_clip=3.,
                                           log_scale_clip_gradient=False,
                                           name=None,
                                           *args,  # pylint: disable=keyword-arg-before-vararg
                                           **kwargs):
  """Build the Masked Autoregressive Density Estimator (Germain et al., 2015).

  This will be wrapped in a make_template to ensure the variables are only
  created once. It takes the input and returns the `loc` ("mu" in [Germain et
  al. (2015)][1]) and `log_scale` ("alpha" in [Germain et al. (2015)][1]) from
  the MADE network.

  Warning: This function uses `masked_dense` to create randomly initialized
  `tf.Variables`. It is presumed that these will be fit, just as you would any
  other neural architecture which uses `tf.layers.dense`.

  #### About Hidden Layers

  Each element of `hidden_layers` should be greater than the `input_depth`
  (i.e., `input_depth = tf.shape(input)[-1]` where `input` is the input to the
  neural network). This is necessary to ensure the autoregressivity property.

  #### About Clipping

  This function also optionally clips the `log_scale` (but possibly not its
  gradient). This is useful because if `log_scale` is too small/large it might
  underflow/overflow making it impossible for the `MaskedAutoregressiveFlow`
  bijector to implement a bijection. Additionally, the `log_scale_clip_gradient`
  `bool` indicates whether the gradient should also be clipped. The default does
  not clip the gradient; this is useful because it still provides gradient
  information (for fitting) yet solves the numerical stability problem. I.e.,
  `log_scale_clip_gradient = False` means
  `grad[exp(clip(x))] = grad[x] exp(clip(x))` rather than the usual
  `grad[clip(x)] exp(clip(x))`.

  Args:
    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. Default: `False`.
    activation: Activation function (callable). Explicitly setting to `None`
      implies a linear activation.
    log_scale_min_clip: `float`-like scalar `Tensor`, or a `Tensor` with the
      same shape as `log_scale`. The minimum value to clip by. Default: -5.
    log_scale_max_clip: `float`-like scalar `Tensor`, or a `Tensor` with the
      same shape as `log_scale`. The maximum value to clip by. Default: 3.
    log_scale_clip_gradient: Python `bool` indicating that the gradient of
      `tf.clip_by_value` should be preserved. Default: `False`.
    name: A name for ops managed by this function. Default:
      "masked_autoregressive_default_template".
    *args: `tf.layers.dense` arguments.
    **kwargs: `tf.layers.dense` keyword arguments.

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

  Raises:
    NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
      graph execution.

  #### References

  [1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE:
       Masked Autoencoder for Distribution Estimation. In _International
       Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509
  """
  name = name or "masked_autoregressive_default_template"
  with tf.compat.v2.name_scope(name):

    def _fn(x):
      """MADE parameterized via `masked_autoregressive_default_template`."""
      # TODO(b/67594795): Better support of dynamic shape.
      input_depth = tf.compat.dimension_value(
          tensorshape_util.with_rank_at_least(x.shape, 1)[-1])
      if input_depth is None:
        raise NotImplementedError(
            "Rightmost dimension must be known prior to graph execution.")
      input_shape = (
          np.int32(tensorshape_util.as_list(x.shape))
          if tensorshape_util.is_fully_defined(x.shape) else tf.shape(input=x))
      if tensorshape_util.rank(x.shape) == 1:
        x = x[tf.newaxis, ...]
      for i, units in enumerate(hidden_layers):
        x = masked_dense(
            inputs=x,
            units=units,
            num_blocks=input_depth,
            exclusive=True if i == 0 else False,
            activation=activation,
            *args,  # pylint: disable=keyword-arg-before-vararg
            **kwargs)
      x = masked_dense(
          inputs=x,
          units=(1 if shift_only else 2) * input_depth,
          num_blocks=input_depth,
          activation=None,
          *args,  # pylint: disable=keyword-arg-before-vararg
          **kwargs)
      if shift_only:
        x = tf.reshape(x, shape=input_shape)
        return x, None
      x = tf.reshape(x, shape=tf.concat([input_shape, [2]], axis=0))
      shift, log_scale = tf.unstack(x, num=2, axis=-1)
      which_clip = (
          tf.clip_by_value
          if log_scale_clip_gradient else clip_by_value_preserve_gradient)
      log_scale = which_clip(log_scale, log_scale_min_clip, log_scale_max_clip)
      return shift, log_scale

    return tf.compat.v1.make_template(name, _fn)


Instruction:
Can you tell what is the following Python 3 function doing
def _create_input_order(input_size, input_order="left-to-right"):
  if isinstance(input_order, six.string_types):
    if input_order == "left-to-right":
      return np.arange(start=1, stop=input_size + 1)
    elif input_order == "right-to-left":
      return np.arange(start=input_size, stop=0, step=-1)
    elif input_order == "random":
      ret = np.arange(start=1, stop=input_size + 1)
      np.random.shuffle(ret)
      return ret
  elif np.all(np.sort(input_order) == np.arange(1, input_size + 1)):
    return np.array(input_order)

  raise ValueError("Invalid input order: '{}'.".format(input_order))

Response: Creates a degree vectors for the input.


Instruction:
Make a summary of the following Python 3 code
def _create_degrees(input_size,
                    hidden_units=None,
                    input_order="left-to-right",
                    hidden_degrees="equal"):
  """Returns a list of degree vectors, one for each input and hidden layer.

  A unit with degree d can only receive input from units with degree < d. Output
  units always have the same degree as their associated input unit.

  Args:
    input_size: Number of inputs.
    hidden_units: list with the number of hidden units per layer. It does not
      include the output layer. Each hidden unit size must be at least the size
      of length (otherwise autoregressivity is not possible).
    input_order: Order of degrees to the input units: 'random', 'left-to-right',
      'right-to-left', or an array of an explicit order. For example,
      'left-to-right' builds an autoregressive model
      p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
    hidden_degrees: Method for assigning degrees to the hidden units:
      'equal', 'random'.  If 'equal', hidden units in each layer are allocated
      equally (up to a remainder term) to each degree.  Default: 'equal'.

  Raises:
    ValueError: invalid input order.
    ValueError: invalid hidden degrees.
  """
  input_order = _create_input_order(input_size, input_order)
  degrees = [input_order]

  if hidden_units is None:
    hidden_units = []

  for units in hidden_units:
    if isinstance(hidden_degrees, six.string_types):
      if hidden_degrees == "random":
        # samples from: [low, high)
        degrees.append(
            np.random.randint(low=min(np.min(degrees[-1]), input_size - 1),
                              high=input_size,
                              size=units))
      elif hidden_degrees == "equal":
        min_degree = min(np.min(degrees[-1]), input_size - 1)
        degrees.append(np.maximum(
            min_degree,
            # Evenly divide the range `[1, input_size - 1]` in to `units + 1`
            # segments, and pick the boundaries between the segments as degrees.
            np.ceil(np.arange(1, units + 1)
                    * (input_size - 1) / float(units + 1)).astype(np.int32)))
    else:
      raise ValueError('Invalid hidden order: "{}".'.format(hidden_degrees))

  return degrees

Response: Creates a list of degree vectors for each input and hidden unit.


Instruction:
Write a Python 3 function that can
return a list of binary mask matrices enforcing autoregressivity.

Response:
def _create_masks(degrees):
  """Returns a list of binary mask matrices enforcing autoregressivity."""
  return [
      # Create input->hidden and hidden->hidden masks.
      inp[:, np.newaxis] <= out
      for inp, out in zip(degrees[:-1], degrees[1:])
  ] + [
      # Create hidden->output mask.
      degrees[-1][:, np.newaxis] < degrees[0]
  ]


Instruction:
Write a Python 3 script for
returning a masked version of the given initializer.

Response:
def _make_masked_initializer(mask, initializer):
  """Returns a masked version of the given initializer."""
  initializer = tf.keras.initializers.get(initializer)
  def masked_initializer(shape, dtype=None, partition_info=None):
    # If no `partition_info` is given, then don't pass it to `initializer`, as
    # `initializer` may be a `tf.compat.v2.initializers.Initializer` (which
    # don't accept a `partition_info` argument).
    if partition_info is None:
      x = initializer(shape, dtype)
    else:
      x = initializer(shape, dtype, partition_info)
    return tf.cast(mask, x.dtype) * x
  return masked_initializer


Instruction:
Write a Python 3 function that can
build the final layer.

Response:
def build(self, input_shape):
    """See tfkl.Layer.build."""
    if self._event_shape is None:
      # `event_shape` wasn't specied at __init__, so infer from `input_shape`.
      self._event_shape = [tf.compat.dimension_value(input_shape[-1])]
      self._event_size = self._event_shape[-1]
      self._event_ndims = len(self._event_shape)
      # Should we throw if input_shape has rank > 2?

    if input_shape[-1] != self._event_shape[-1]:
      raise ValueError("Invalid final dimension of `input_shape`. "
                       "Expected `{!r}`, but got `{!r}`".format(
                           self._event_shape[-1], input_shape[-1]))

    # Construct the masks.
    self._input_order = _create_input_order(
        self._event_size, self._input_order_param)
    self._masks = _create_masks(_create_degrees(
        input_size=self._event_size,
        hidden_units=self._hidden_units,
        input_order=self._input_order,
        hidden_degrees=self._hidden_degrees))

    # In the final layer, we will produce `self._params` outputs for each of the
    # `self._event_size` inputs to `AutoregressiveLayer`.  But `masks[-1]` has
    # shape `[self._hidden_units[-1], self._event_size]`.  Thus, we need to
    # expand the mask to `[hidden_units[-1], event_size * self._params]` such
    # that all units for the same input are masked identically.  In particular,
    # we tile the mask so the j-th element of `tf.unstack(output, axis=-1)` is a
    # tensor of the j-th parameter/unit for each input.
    #
    # NOTE: Other orderings of the output could be faster -- should benchmark.
    self._masks[-1] = np.reshape(
        np.tile(self._masks[-1][..., tf.newaxis], [1, 1, self._params]),
        [self._masks[-1].shape[0], self._event_size * self._params])

    self._network = tf.keras.Sequential([
        # Starting this model with an `InputLayer` ensures that Keras will build
        # and propagate our `dtype` to each layer we add.
        tf.keras.layers.InputLayer((self._event_size,), dtype=self.dtype)
    ])

    # Input-to-hidden, hidden-to-hidden, and hidden-to-output layers:
    #  [..., self._event_size] -> [..., self._hidden_units[0]].
    #  [..., self._hidden_units[k-1]] -> [..., self._hidden_units[k]].
    #  [..., self._hidden_units[-1]] -> [..., event_size * self._params].
    layer_output_sizes = self._hidden_units + [self._event_size * self._params]
    for k in range(len(self._masks)):
      self._network.add(tf.keras.layers.Dense(
          layer_output_sizes[k],
          kernel_initializer=_make_masked_initializer(
              self._masks[k], self._kernel_initializer),
          kernel_constraint=_make_masked_constraint(self._masks[k]),
          activation=self._activation if k + 1 < len(self._masks) else None,
          use_bias=self._use_bias,
          **self._kwargs))

    # Record that the layer has been built.
    super(AutoregressiveLayer, self).build(input_shape)


Instruction:
Given the following Python 3 function, write the documentation
def call(self, x):
    with tf.compat.v2.name_scope(self.name or "AutoregressiveLayer_call"):
      x = tf.convert_to_tensor(value=x, dtype=self.dtype, name="x")
      input_shape = tf.shape(input=x)
      # TODO(b/67594795): Better support for dynamic shapes.
      if tensorshape_util.rank(x.shape) == 1:
        x = x[tf.newaxis, ...]
      return tf.reshape(self._network(x),
                        tf.concat([input_shape, [self._params]], axis=0))

Response: See tf. Layer. call.


Instruction:
Create a Python 3 function for
drawing samples from a multinomial dataset.

Response:
def draw_sample(num_samples, num_classes, logits, num_trials, dtype, seed):
  """Sample a multinomial.

  The batch shape is given by broadcasting num_trials with
  remove_last_dimension(logits).

  Args:
    num_samples: Python int or singleton integer Tensor: number of multinomial
      samples to draw.
    num_classes: Python int or singleton integer Tensor: number of classes.
    logits: Floating Tensor with last dimension k, of (unnormalized) logit
      probabilities per class.
    num_trials: Tensor of number of categorical trials each multinomial consists
      of.  num_trials[..., tf.newaxis] must broadcast with logits.
    dtype: dtype at which to emit samples.
    seed: Random seed.

  Returns:
    samples: Tensor of given dtype and shape [n] + batch_shape + [k].
  """
  with tf.name_scope("multinomial.draw_sample"):
    # broadcast the num_trials and logits to same shape
    num_trials = tf.ones_like(
        logits[..., 0], dtype=num_trials.dtype) * num_trials
    logits = tf.ones_like(
        num_trials[..., tf.newaxis], dtype=logits.dtype) * logits

    # flatten the total_count and logits
    # flat_logits has shape [B1B2...Bm, num_classes]
    flat_logits = tf.reshape(logits, [-1, num_classes])
    flat_num_trials = num_samples * tf.reshape(num_trials, [-1])  # [B1B2...Bm]

    # Computes each logits and num_trials situation by map_fn.

    # Using just one batch tf.random.categorical call doesn't work because that
    # requires num_trials to be the same across all members of the batch of
    # logits.  This restriction makes sense for tf.random.categorical because
    # for it, num_trials is part of the returned shape.  However, the
    # multinomial sampler does not need that restriction, because it sums out
    # exactly that dimension.

    # One possibility would be to draw a batch categorical whose sample count is
    # max(num_trials) and mask out the excess ones.  However, if the elements of
    # num_trials vary widely, this can be wasteful of memory.

    # TODO(b/123763054, b/112152209): Revisit the possibility of writing this
    # with a batch categorical followed by batch unsorted_segment_sum, once both
    # of those work and are memory-efficient enough.
    def _sample_one_batch_member(args):
      logits, num_cat_samples = args[0], args[1]  # [K], []
      # x has shape [1, num_cat_samples = num_samples * num_trials]
      x = tf.random.categorical(
          logits[tf.newaxis, ...], num_cat_samples, seed=seed)
      x = tf.reshape(x, shape=[num_samples, -1])  # [num_samples, num_trials]
      x = tf.one_hot(
          x, depth=num_classes)  # [num_samples, num_trials, num_classes]
      x = tf.reduce_sum(input_tensor=x, axis=-2)  # [num_samples, num_classes]
      return tf.cast(x, dtype=dtype)

    x = tf.map_fn(
        _sample_one_batch_member, [flat_logits, flat_num_trials],
        dtype=dtype)  # [B1B2...Bm, num_samples, num_classes]

    # reshape the results to proper shape
    x = tf.transpose(a=x, perm=[1, 0, 2])
    final_shape = tf.concat([[num_samples],
                             tf.shape(input=num_trials), [num_classes]],
                            axis=0)
    x = tf.reshape(x, final_shape)

    return x


Instruction:
How would you code a function in Python 3 to
build a zero - dimensional MVNDiag object.

Response:
def _zero_dimensional_mvndiag(dtype):
  """Build a zero-dimensional MVNDiag object."""
  dummy_mvndiag = tfd.MultivariateNormalDiag(
      scale_diag=tf.ones([0], dtype=dtype))
  dummy_mvndiag.covariance = lambda: dummy_mvndiag.variance()[..., tf.newaxis]
  return dummy_mvndiag


Instruction:
Write a Python 3 function that can
build an observation_noise_fn that observes a Tensor timeseries.

Response:
def _observe_timeseries_fn(timeseries):
  """Build an observation_noise_fn that observes a Tensor timeseries."""
  def observation_noise_fn(t):
    current_slice = timeseries[..., t, :]
    return tfd.MultivariateNormalDiag(
        loc=current_slice,
        scale_diag=tf.zeros_like(current_slice))
  return observation_noise_fn


Instruction:
Can you generate the documentation for the following Python 3 function
def params_to_weights(self,
                        global_scale_variance,
                        global_scale_noncentered,
                        local_scale_variances,
                        local_scales_noncentered,
                        weights_noncentered):
    """Build regression weights from model parameters."""
    global_scale = (global_scale_noncentered *
                    tf.sqrt(global_scale_variance) *
                    self.weights_prior_scale)

    local_scales = local_scales_noncentered * tf.sqrt(local_scale_variances)
    return weights_noncentered * local_scales * global_scale[..., tf.newaxis]

Response: Build regression weights from model parameters.


Instruction:
Can you write a function in Python 3 where it
computes the number of edges on longest path from node to root.

Response:
def _depth(g):
  """Computes the number of edges on longest path from node to root."""
  def _explore(v):
    if v.depth < 0:
      v.depth = ((1 + max([-1] + [_explore(annotated_graph[u])
                                  for u in v.parents]))
                 if v.parents else 0)
    return v.depth
  annotated_graph = {k: _Node(k, v) for k, v in g.items()}
  for v in annotated_graph.values():
    _explore(v)
  return annotated_graph


Instruction:
How would you implement a function in Python 3 that
returns a tuple of str tuple - str pairs representing resolved & sorted DAG.

Response:
def _best_order(g):
  """Creates tuple of str tuple-str pairs representing resolved & sorted DAG."""
  def _explore(u):
    """Recursive function to ascend up through unvisited dependencies."""
    if u.depth < 0:
      return  # Already visited.
    if not u.parents:
      result.append((u.name, u.parents))
      u.depth = -1  # Mark visited.
      return
    b = (u.name, [])
    result.append(b)
    u.depth = -1  # Mark visited.
    d = 0
    for v in sorted((g.get(p) for p in u.parents), key=lambda v: v.depth):
      n0 = len(result)
      _explore(v)
      n1 = len(result)
      b[1].extend(['_']*d + [v.name])
      d = n1 - n0 - 1
  g = _depth(g)
  result = []
  for u in sorted(g.values(), key=lambda v: v.depth, reverse=True):
    _explore(u)
  return tuple(reversed(result))


Instruction:
Can you tell what is the following Python 3 function doing
def _prob_chain_rule_flatten(named_makers):
  def _make(dist_fn, args):
    if args is None:
      return lambda *_: dist_fn
    if not args:
      return lambda *_: dist_fn()
    def _fn(*xs):
      kwargs = dict(zip(args, reversed(xs[-len(args):])))
      kwargs.pop('_', None)
      return dist_fn(**kwargs)
    return _fn
  named_makers = _convert_to_dict(named_makers)
  g = {k: (None if distribution_util.is_distribution_instance(v)
           else joint_distribution_sequential._get_required_args(v))  # pylint: disable=protected-access
       for k, v in named_makers.items()}
  g = _best_order(g)
  dist_fn_name, dist_fn_args = zip(*g)
  dist_fn_args = tuple(None if a is None else tuple(a) for a in dist_fn_args)
  dist_fn_wrapped = tuple(_make(named_makers[name], parents)
                          for (name, parents) in g)
  dist_fn = tuple(named_makers.get(n) for n in dist_fn_name)
  return dist_fn, dist_fn_wrapped, dist_fn_args, dist_fn_name

Response: Creates lists of callables suitable for JDSeq.


Instruction:
Given the following Python 3 function, write the documentation
def _build(self, model):
    if not _is_dict_like(model):
      raise TypeError('`model` must be convertible to `dict` (saw: {}).'.format(
          type(model).__name__))
    [
        self._dist_fn,
        self._dist_fn_wrapped,
        self._dist_fn_args,
        self._dist_fn_name,  # JointDistributionSequential doesn't have this.
    ] = _prob_chain_rule_flatten(model)

Response: Creates the distribution function and distribution function list from the given model.


Instruction:
Explain what the following Python 3 code does
def variational_loss(self,
                       observations,
                       observation_index_points=None,
                       kl_weight=1.,
                       name='variational_loss'):
    """Variational loss for the VGP.

    Given `observations` and `observation_index_points`, compute the
    negative variational lower bound as specified in [Hensman, 2013][1].

    Args:
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `observation_index_points`, etc.).
      observation_index_points: `float` `Tensor` representing finite (batch of)
        vector(s) of points where observations are defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below). If
        set to `None` uses `index_points` as the origin for observations.
        Default value: None.
      kl_weight: Amount by which to scale the KL divergence loss between prior
        and posterior.
        Default value: 1.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "GaussianProcess".
    Returns:
      loss: Scalar tensor representing the negative variational lower bound.
        Can be directly used in a `tf.Optimizer`.
    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.

    #### References

    [1]: Hensman, J., Lawrence, N. "Gaussian Processes for Big Data", 2013
         https://arxiv.org/abs/1309.6835
    """

    with tf.name_scope(name or 'variational_gp_loss'):
      if observation_index_points is None:
        observation_index_points = self._index_points
      observation_index_points = tf.convert_to_tensor(
          value=observation_index_points, dtype=self._dtype,
          name='observation_index_points')
      observations = tf.convert_to_tensor(
          value=observations, dtype=self._dtype, name='observations')
      kl_weight = tf.convert_to_tensor(
          value=kl_weight, dtype=self._dtype,
          name='kl_weight')

      # The variational loss is a negative ELBO. The ELBO can be broken down
      # into three terms:
      #  1. a likelihood term
      #  2. a trace term arising from the covariance of the posterior predictive

      kzx = self.kernel.matrix(self._inducing_index_points,
                               observation_index_points)

      kzx_linop = tf.linalg.LinearOperatorFullMatrix(kzx)
      loc = (self._mean_fn(observation_index_points) +
             kzx_linop.matvec(self._kzz_inv_varloc, adjoint=True))

      likelihood = independent.Independent(
          normal.Normal(
              loc=loc,
              scale=tf.sqrt(self._observation_noise_variance + self._jitter),
              name='NormalLikelihood'),
          reinterpreted_batch_ndims=1)
      obs_ll = likelihood.log_prob(observations)

      chol_kzz_linop = tf.linalg.LinearOperatorLowerTriangular(self._chol_kzz)
      chol_kzz_inv_kzx = chol_kzz_linop.solve(kzx)
      kzz_inv_kzx = chol_kzz_linop.solve(chol_kzz_inv_kzx, adjoint=True)

      kxx_diag = tf.linalg.diag_part(
          self.kernel.matrix(
              observation_index_points, observation_index_points))
      ktilde_trace_term = (
          tf.reduce_sum(input_tensor=kxx_diag, axis=-1) -
          tf.reduce_sum(input_tensor=chol_kzz_inv_kzx ** 2, axis=[-2, -1]))

      # Tr(SB)
      # where S = A A.T, A = variational_inducing_observations_scale
      # and B = Kzz^-1 Kzx Kzx.T Kzz^-1
      #
      # Now Tr(SB) = Tr(A A.T Kzz^-1 Kzx Kzx.T Kzz^-1)
      #            = Tr(A.T Kzz^-1 Kzx Kzx.T Kzz^-1 A)
      #            = sum_ij (A.T Kzz^-1 Kzx)_{ij}^2
      other_trace_term = tf.reduce_sum(
          input_tensor=(
              self._variational_inducing_observations_posterior.scale.matmul(
                  kzz_inv_kzx) ** 2),
          axis=[-2, -1])

      trace_term = (.5 * (ktilde_trace_term + other_trace_term) /
                    self._observation_noise_variance)

      inducing_prior = gaussian_process.GaussianProcess(
          kernel=self._kernel,
          mean_fn=self._mean_fn,
          index_points=self._inducing_index_points,
          observation_noise_variance=self._observation_noise_variance)

      kl_term = kl_weight * kullback_leibler.kl_divergence(
          self._variational_inducing_observations_posterior,
          inducing_prior)

      lower_bound = (obs_ll - trace_term - kl_term)

      return -tf.reduce_mean(input_tensor=lower_bound)

Response: Returns the variational loss for the VGP.