Instruction:
Create a Python 3 function to
return a cloned Linear module.

Response:
def clone(self, name=None):
    """Returns a cloned `Linear` module.

    Args:
      name: Optional string assigning name of cloned module. The default name
          is constructed by appending "_clone" to `self.module_name`.

    Returns:
      Cloned `Linear` module.
    """
    if name is None:
      name = self.module_name + "_clone"
    return Linear(output_size=self.output_size,
                  use_bias=self._use_bias,
                  initializers=self._initializers,
                  partitioners=self._partitioners,
                  regularizers=self._regularizers,
                  name=name)


Instruction:
Given the following Python 3 function, write the documentation
def _build(self, inputs_list):
    outputs = []
    for idx, tensor in enumerate(inputs_list):
      outputs.append(
          Linear(
              self._output_size,
              initializers=self._initializers,
              partitioners=self._partitioners,
              regularizers=self._regularizers,
              # Since we are interpreting this as 'one big linear', we only need
              # one bias.
              use_bias=(idx == 0 and self._use_bias))(tensor))
    return tf.add_n(outputs)

Response: Connects the module into the graph.


Instruction:
Write a Python 3 script for
connecting the Add module into the graph with input Tensor inputs.

Response:
def _build(self, inputs, multiplier=1):
    """Connects the Add module into the graph, with input Tensor `inputs`.

    Args:
      inputs: A Tensor of size `[batch_size, input_size1, ...]`.
      multiplier: A scalar or Tensor which the bias term is multiplied by
        before adding it to `inputs`. Anything which works in the expression
        `bias * multiplier` is acceptable here. This may be useful if you want
        to add a bias in one place and subtract the same bias in another place
        via `multiplier=-1`.

    Returns:
      A Tensor of size `[batch_size, input_size1, ...]`.

    Raises:
      base.IncompatibleShapeError: If the input is not a >= 2D `Tensor`.
      base.IncompatibleShapeError: If connecting the module into the graph
          any time after the first time, and the inferred size of the input does
          not match previous invocations.
      base.IncompatibleShapeError: If the `output_shape` has been specified
          but it does not match the input_shape`.
      base.ParentNotBuiltError: If the module is a transposed and the original
          untransposed module has not been built.
    """
    input_shape = tuple(inputs.get_shape().as_list())
    bias_shape = calculate_bias_shape(input_shape, self._bias_dims)

    # Check always contains minibatched input.
    if len(input_shape) < 2:
      raise base.IncompatibleShapeError(
          "Rank of input shape must be >=2 not: {}.".format(len(input_shape)))

    # Check previous input size is same as new input size.
    if (self._input_shape is not None and
        input_shape[1:] != self._input_shape[1:]):
      raise base.IncompatibleShapeError("Input shape has changed.")

    # If transposed, make sure that the original Module is built.
    if callable(self._output_shape):
      self._output_shape = self._output_shape()
      if self._output_shape is None:
        raise base.ParentNotBuiltError(
            "Build the original untransposed module before building this one.")

    # If output_shape specified, check that it matches input_shape.
    if (self._output_shape is not None and
        self._output_shape[1:] != input_shape[1:]):
      raise base.IncompatibleShapeError(
          "Input shape must be {} not: {}.".format(self._output_shape,
                                                   input_shape[1]))

    self._input_shape = input_shape
    dtype = inputs.dtype

    if "b" not in self._initializers:
      self._initializers["b"] = create_bias_initializer(bias_shape, dtype)

    self._b = tf.get_variable(
        "b",
        shape=bias_shape,
        dtype=dtype,
        initializer=self._initializers["b"],
        partitioner=self._partitioners.get("b", None),
        regularizer=self._regularizers.get("b", None))

    bias = self._b
    if multiplier != 1:
      bias = bias * multiplier  # pylint: disable=g-no-augmented-assignment
    outputs = inputs + bias
    return outputs


Instruction:
Write a Python 3 script for
returning a transpose of the add bias module.

Response:
def transpose(self, name=None):
    """Returns transposed `AddBias` module.

    Args:
      name: Optional string assigning name of transpose module. The default name
          is constructed by appending "_transpose" to `self.module_name`.

    Returns:
      Transposed `AddBias` module.
    """

    if name is None:
      name = self.module_name + "_transpose"
    return AddBias(output_shape=lambda: self._input_shape,
                   bias_dims=self._bias_dims,
                   initializers=self._initializers,
                   regularizers=self._regularizers,
                   name=name)


Instruction:
Explain what the following Python 3 code does
def _infer_shape(self, dimensions):
    # Size of input
    n = np.prod(dimensions)
    # Size of output where defined
    m = np.prod(abs(np.array(self._shape)))
    # Replace wildcard
    v = np.array(self._shape)
    v[v == -1] = n // m
    return tuple(v)

Response: Infer the output shape vector given the input dimensions.


Instruction:
Can you generate the documentation for the following Python 3 function
def _build(self, inputs):
    full_input_shape = inputs.get_shape().as_list()
    if len(full_input_shape) < self._preserve_dims:
      raise ValueError("Input tensor has {} dimensions, should have at least "
                       "as many as preserve_dims={}".format(
                           len(full_input_shape),
                           self._preserve_dims))
    self._input_shape = full_input_shape[self._preserve_dims:]

    if callable(self._shape):
      self._shape = tuple(self._shape())

    # Special-case of 1 non-preserved dimension, where no reshape is necessary.
    # This is useful if the non-preserved dimension of `inputs` is unknown
    # at build time.
    if len(self._input_shape) == 1 and len(self._shape) == 1:
      if self._shape[0] == -1 or self._shape[0] == self._input_shape[0]:
        return inputs
      else:
        if self._input_shape[0] is None:
          raise ValueError("Unknown non-preserved dimensions are not allowed "
                           "in the input to BatchReshape unless it is only one "
                           "and the desired shape is (-1,).")
        else:
          raise ValueError("Output shape is incompatible with input shape")

    if not all([isinstance(x, numbers.Integral) and (x > 0 or x == -1)
                for x in self._shape]):
      raise ValueError(
          "Desired shape can only contain positive integral numbers "
          "and the wildcard -1. Given shape {}".format(self._shape))

    if self._shape.count(-1) > 1:
      raise ValueError(
          "Wildcard -1 can appear only once in desired output shape. "
          "Given shape {}".format(self._shape))

    preserved_shape = tf.shape(inputs)[:self._preserve_dims]
    # Slicing the shape tensor loses information, we keep it in a list.
    preserved_shape_list = inputs.get_shape()[:self._preserve_dims]

    # Except in the case above where no reshape is needed, we do not allow
    # unknown non-preserved dimensions in the input.
    if None in self._input_shape:
      raise ValueError("Unknown non-preserved dimensions are not allowed in "
                       "the input to BatchReshape unless it is only one and the"
                       " desired shape is (-1,). The offending non-preserved "
                       "input shape is {}".format(self._input_shape))
    if self._shape.count(-1) > 0:
      trailing_shape = self._infer_shape(self._input_shape)
    else:
      trailing_shape = self._shape

    if np.prod(self._input_shape) != np.prod(trailing_shape):
      raise ValueError("Output shape is incompatible with input shape")

    shape = tf.concat([preserved_shape, trailing_shape], 0)
    output = tf.reshape(inputs, shape)

    # Include shape information that was lost when we sliced the shape tensor.
    shape_list = preserved_shape_list.concatenate(trailing_shape)
    output.set_shape(shape_list)

    return output

Response: Connects the module into the graph with input Tensor inputs.


Instruction:
Write a Python 3 function that can
return transpose batch reshape.

Response:
def transpose(self, name=None):
    """Returns transpose batch reshape."""
    if name is None:
      name = self.module_name + "_transpose"
    return BatchReshape(shape=lambda: self.input_shape,
                        preserve_dims=self._preserve_dims,
                        name=name)


Instruction:
Can you generate a brief explanation for the following Python 3 code
def _build(self):
    if "w" not in self._initializers:
      stddev = 1 / math.sqrt(np.prod(self._shape))
      self._initializers["w"] = tf.truncated_normal_initializer(stddev=stddev)

    self._w = tf.get_variable("w",
                              shape=self._shape,
                              dtype=self._dtype,
                              initializer=self._initializers["w"],
                              partitioner=self._partitioners.get("w", None),
                              regularizer=self._regularizers.get("w", None))
    return self._w

Response: Connects the TrainableTensor module into the graph.


Instruction:
How would you implement a function in Python 3 that
connects the BatchApply module into the graph.

Response:
def _build(self, *args, **kwargs):
    """Connects the BatchApply module into the graph.

    Args:
      *args: a Tensor or a nested list or dictionary of Tensors. The input
          tensors will have their first dimensions merged, then an op or a
          module will be called on the input. The first dimension of the output
          tensor(s) will be split again based on the leading dimensions of the
          first input tensor.
      **kwargs: Dictionary of named arguments; used in the same way as `*args`.

    Returns:
      A Tensor or nested list or dictionary of Tensors as a result of applying
      the process above. ("None" return values are also supported.)
    """
    flattened = nest.flatten([args, kwargs])
    merged_flattened = [
        merge_leading_dims(inp, self._n_dims) if inp is not None else None
        for inp in flattened]
    merged_args, merged_kwargs = nest.pack_sequence_as([args, kwargs],
                                                       merged_flattened)

    results = self._module(*merged_args, **merged_kwargs)

    # Unmerging takes the sizes of the leading dimensions from an input example
    # with equal shape for the leading `n_dims` dimensions. Typically this is
    # the first input.
    example_input = tf.convert_to_tensor(flattened[self._input_example_index])
    def _split_to_original_leading_dims(result):
      if result is None:
        return None
      else:
        return split_leading_dim(result, example_input, self._n_dims)

    flat_results = nest.flatten(results)
    flat_unmerged_results = [_split_to_original_leading_dims(result)
                             for result in flat_results]
    return nest.pack_sequence_as(results, flat_unmerged_results)


Instruction:
Write a Python 3 script to
connect the SliceByDim module into the graph.

Response:
def _build(self, inputs):
    """Connects the SliceByDim module into the graph.

    Args:
      inputs: `Tensor` to slice. Its rank must be greater than the maximum
          dimension specified in `dims` (plus one as python is 0 indexed).

    Returns:
      The sliced tensor.

    Raises:
      ValueError: If `inputs` tensor has insufficient rank.
    """
    shape_inputs = inputs.get_shape().as_list()
    rank = len(shape_inputs)

    # Checks that the rank of the tensor.
    max_dim = np.max(self._dims) + 1
    if rank < max_dim:
      raise ValueError("Rank of inputs must be at least {}.".format(max_dim))

    # Builds default lists for begin and size to pass to `tf.slice`.
    full_begin = [0] * rank
    full_size = [-1] * rank

    # Updates lists with what the user provided.
    for dim, begin, size in zip(self._dims, self._begin, self._size):
      full_begin[dim] = begin
      full_size[dim] = size

    return tf.slice(inputs, begin=full_begin, size=full_size)


Instruction:
Write a Python 3 function that can
connect the TileByDim module into the graph.

Response:
def _build(self, inputs):
    """Connects the `TileByDim` module into the graph.

    Args:
      inputs: `Tensor` to tile.

    Returns:
      The tiled tensor.
    """
    shape_inputs = inputs.get_shape().as_list()
    rank = len(shape_inputs)

    # Builds default lists for multiples to pass to `tf.tile`.
    full_multiples = [1] * rank

    # Updates lists with what the user provided.
    for dim, multiple in zip(self._dims, self._multiples):
      full_multiples[dim] = multiple

    return tf.tile(inputs, multiples=full_multiples)


Instruction:
Write a Python 3 function that can
connect the MergeDims module into the graph.

Response:
def _build(self, inputs):
    """Connects the MergeDims module into the graph.

    Args:
      inputs: Tensor or a nested list of Tensors to merge. Its rank must be
          greater than or equal to `start` + `size`.

    Returns:
      The merged Tensor or a nested list of merged Tensors.

    Raises:
      ValueError: If any of the `inputs` tensors has insufficient rank.
    """
    if nest.is_sequence(inputs):
      merged_tensors = [self._merge(tensor) for tensor in nest.flatten(inputs)]
      return nest.pack_sequence_as(inputs, merged_tensors)

    # inputs is a single tf.Tensor
    return self._merge(inputs)


Instruction:
Can you generate the documentation for the following Python 3 function
def initial_state(self, batch_size, trainable=False):
    init_state = tf.eye(self._mem_slots, batch_shape=[batch_size])

    # Pad the matrix with zeros.
    if self._mem_size > self._mem_slots:
      difference = self._mem_size - self._mem_slots
      pad = tf.zeros((batch_size, self._mem_slots, difference))
      init_state = tf.concat([init_state, pad], -1)
    # Truncation. Take the first `self._mem_size` components.
    elif self._mem_size < self._mem_slots:
      init_state = init_state[:, :, :self._mem_size]
    return init_state

Response: Creates the initial state of the memory.


Instruction:
How would you code a function in Python 3 to
perform multi - head attention from Attention is All You Need.

Response:
def _multihead_attention(self, memory):
    """Perform multi-head attention from 'Attention is All You Need'.

    Implementation of the attention mechanism from
    https://arxiv.org/abs/1706.03762.

    Args:
      memory: Memory tensor to perform attention on.

    Returns:
      new_memory: New memory tensor.
    """
    key_size = self._key_size
    value_size = self._head_size

    qkv_size = 2 * key_size + value_size
    total_size = qkv_size * self._num_heads  # Denote as F.
    qkv = basic.BatchApply(basic.Linear(total_size))(memory)
    qkv = basic.BatchApply(layer_norm.LayerNorm())(qkv)

    mem_slots = memory.get_shape().as_list()[1]  # Denoted as N.

    # [B, N, F] -> [B, N, H, F/H]
    qkv_reshape = basic.BatchReshape([mem_slots, self._num_heads,
                                      qkv_size])(qkv)

    # [B, N, H, F/H] -> [B, H, N, F/H]
    qkv_transpose = tf.transpose(qkv_reshape, [0, 2, 1, 3])
    q, k, v = tf.split(qkv_transpose, [key_size, key_size, value_size], -1)

    q *= key_size ** -0.5
    dot_product = tf.matmul(q, k, transpose_b=True)  # [B, H, N, N]
    weights = tf.nn.softmax(dot_product)

    output = tf.matmul(weights, v)  # [B, H, N, V]

    # [B, H, N, V] -> [B, N, H, V]
    output_transpose = tf.transpose(output, [0, 2, 1, 3])

    # [B, N, H, V] -> [B, N, H * V]
    new_memory = basic.BatchFlatten(preserve_dims=2)(output_transpose)
    return new_memory


Instruction:
Write a Python 3 function that can
create input and forget gates for this step using inputs and memory.

Response:
def _create_gates(self, inputs, memory):
    """Create input and forget gates for this step using `inputs` and `memory`.

    Args:
      inputs: Tensor input.
      memory: The current state of memory.

    Returns:
      input_gate: A LSTM-like insert gate.
      forget_gate: A LSTM-like forget gate.
    """
    # We'll create the input and forget gates at once. Hence, calculate double
    # the gate size.
    num_gates = 2 * self._calculate_gate_size()

    memory = tf.tanh(memory)
    inputs = basic.BatchFlatten()(inputs)
    gate_inputs = basic.BatchApply(basic.Linear(num_gates), n_dims=1)(inputs)
    gate_inputs = tf.expand_dims(gate_inputs, axis=1)
    gate_memory = basic.BatchApply(basic.Linear(num_gates))(memory)
    gates = tf.split(gate_memory + gate_inputs, num_or_size_splits=2, axis=2)
    input_gate, forget_gate = gates

    input_gate = tf.sigmoid(input_gate + self._input_bias)
    forget_gate = tf.sigmoid(forget_gate + self._forget_bias)

    return input_gate, forget_gate


Instruction:
Write a Python 3 script for
performing multiheaded attention over memory.

Response:
def _attend_over_memory(self, memory):
    """Perform multiheaded attention over `memory`.

    Args:
      memory: Current relational memory.

    Returns:
      The attended-over memory.
    """
    attention_mlp = basic.BatchApply(
        mlp.MLP([self._mem_size] * self._attention_mlp_layers))
    for _ in range(self._num_blocks):
      attended_memory = self._multihead_attention(memory)

      # Add a skip connection to the multiheaded attention's input.
      memory = basic.BatchApply(layer_norm.LayerNorm())(
          memory + attended_memory)

      # Add a skip connection to the attention_mlp's input.
      memory = basic.BatchApply(layer_norm.LayerNorm())(
          attention_mlp(memory) + memory)

    return memory


Instruction:
Given the following Python 3 function, write the documentation
def _build(self, inputs, memory, treat_input_as_matrix=False):
    if treat_input_as_matrix:
      inputs = basic.BatchFlatten(preserve_dims=2)(inputs)
      inputs_reshape = basic.BatchApply(
          basic.Linear(self._mem_size), n_dims=2)(inputs)
    else:
      inputs = basic.BatchFlatten()(inputs)
      inputs = basic.Linear(self._mem_size)(inputs)
      inputs_reshape = tf.expand_dims(inputs, 1)

    memory_plus_input = tf.concat([memory, inputs_reshape], axis=1)
    next_memory = self._attend_over_memory(memory_plus_input)

    n = inputs_reshape.get_shape().as_list()[1]
    next_memory = next_memory[:, :-n, :]

    if self._gate_style == 'unit' or self._gate_style == 'memory':
      self._input_gate, self._forget_gate = self._create_gates(
          inputs_reshape, memory)
      next_memory = self._input_gate * tf.tanh(next_memory)
      next_memory += self._forget_gate * memory

    output = basic.BatchFlatten()(next_memory)
    return output, next_memory

Response: Builds the graph for the time step.


Instruction:
Here you have a function in Python 3, explain what it does
def train_and_eval(train_batch_size, test_batch_size, num_hidden, learning_rate,
                   num_train_steps, report_every, test_every):
  """Creates a basic MNIST model using Sonnet, then trains and evaluates it."""

  data_dict = dataset_mnist.get_data("mnist", train_batch_size, test_batch_size)
  train_data = data_dict["train_iterator"]
  test_data = data_dict["test_iterator"]

  # Sonnet separates the configuration of a model from its attachment into the
  # graph. Here we configure the shape of the model, but this call does not
  # place any ops into the graph.
  mlp = snt.nets.MLP([num_hidden, data_dict["num_classes"]])

  train_images, train_labels = train_data.get_next()
  test_images, test_labels = test_data.get_next()

  # Flatten images to pass to model.
  train_images = snt.BatchFlatten()(train_images)
  test_images = snt.BatchFlatten()(test_images)

  # Call our model, which creates it in the graph. Our build function
  # is parameterized by the source of images, and here we connect the model to
  # the training images.
  train_logits = mlp(train_images)

  # Training loss and optimizer.
  loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=train_labels, logits=train_logits)
  loss_avg = tf.reduce_mean(loss)
  optimizer = tf.train.GradientDescentOptimizer(learning_rate)
  optimizer_step = optimizer.minimize(loss_avg)

  # As before, we make a second instance of our model in the graph, which shares
  # its parameters with the first instance of the model. The Sonnet Module code
  # takes care of the variable sharing for us: because we are calling the same
  # instance of Model, we will automatically reference the same, shared
  # variables.
  test_logits = mlp(test_images)
  test_classes = tf.nn.softmax(test_logits)
  test_correct = tf.nn.in_top_k(test_classes, test_labels, k=1)

  with tf.train.SingularMonitoredSession() as sess:

    for step_idx in range(num_train_steps):
      current_loss, _ = sess.run([loss_avg, optimizer_step])
      if step_idx % report_every == 0:
        tf.logging.info("Step: %4d of %d - loss: %.02f.",
                        step_idx + 1, num_train_steps, current_loss)
      if step_idx % test_every == 0:
        sess.run(test_data.initializer)
        current_correct = sess.run(test_correct)
        correct_count = np.count_nonzero(current_correct)
        tf.logging.info("Test: %d of %d correct.",
                        correct_count, test_batch_size)

Response: Creates a basic MNIST model using Sonnet.


Instruction:
Implement a function in Python 3 to
return an initial learnable state.

Response:
def _single_learnable_state(state, state_id=0, learnable=True):
  """Returns an initial (maybe learnable) state.

  This function does not create any variable scopes, and it should be called
  from a Sonnet module. This function also makes sure that all the rows of its
  `state` argument have the same value.

  Args:
    state: initial value of the initial state. It should be a tensor of at least
      two dimensions, of which the first dimension corresponds to the
      batch_size dimension. All rows of such tensor should have the same value.
    state_id: integer that uniquely identifies this state.
    learnable: boolean that indicates whether the state is learnable.

  Returns:
    The initial learnable state `Tensor`.
  """
  unpacked_state = tf.unstack(state)
  # Assert that all rows have the same values.
  assert_rows_equal = [tf.assert_equal(s, unpacked_state[0])
                       for s in unpacked_state]

  # We wish to have all the graph assertions in the graph's critical path,
  # so we include them even if the initial state is left unmodified (i.e. when
  # the state is not learnable).
  # Note: All these assertions will be run every time that data flows
  # through the graph. At that point, the control_dependencies context manager
  # makes sure that such assertions are run, and will raise an exception if any
  # fails.
  with tf.control_dependencies(assert_rows_equal):
    if not learnable:
      return state
    else:
      state_shape = state.get_shape()
      state_shape.assert_is_fully_defined()
      state_shape_list = state_shape.as_list()
      batch_size, trailing_shape = state_shape_list[0], state_shape_list[1:]

      initial_value = tf.reshape(unpacked_state[0], [1] + trailing_shape)
      initial_state_variable = tf.get_variable(
          "initial_state_%d" % state_id, dtype=initial_value.dtype,
          initializer=initial_value)

      trailing_size_repeat = [1] * len(trailing_shape)
      return tf.tile(initial_state_variable,
                     tf.constant([batch_size] + trailing_size_repeat))


Instruction:
Can you write a function in Python 3 where it
returns an initial state consisting of trainable variables.

Response:
def trainable_initial_state(batch_size, state_size, dtype, initializers=None,
                            regularizers=None, name=None):
  """Creates an initial state consisting of trainable variables.

  The trainable variables are created with the same shapes as the elements of
  `state_size` and are tiled to produce an initial state.

  Args:
    batch_size: An int, or scalar int32 Tensor representing the batch size.
    state_size: A `TensorShape` or nested tuple of `TensorShape`s to use for the
        shape of the trainable variables.
    dtype: The data type used to create the variables and thus initial state.
    initializers: An optional container of the same structure as `state_size`
        containing initializers for the variables.
    regularizers: An optional container of the same structure as `state_size`
        containing regularizers for the variables.
    name: optional string used to prefix the initial state variable names.

  Returns:
    A `Tensor` or nested tuple of `Tensor`s with the same size and structure
    as `state_size`, where each `Tensor` is a tiled trainable `Variable`.

  Raises:
    ValueError: if the user passes initializers that are not functions.
    ValueError: if the user passes regularizers that are not functions.
  """
  flat_state_size = nest.flatten(state_size)

  if not initializers:
    flat_initializer = tuple(tf.zeros_initializer() for _ in flat_state_size)
  else:
    nest.assert_same_structure(initializers, state_size)
    flat_initializer = nest.flatten(initializers)
    if not all([callable(init) for init in flat_initializer]):
      raise ValueError("Not all the passed initializers are callable objects.")

  if not regularizers:
    flat_regularizer = tuple({} for _ in flat_state_size)
  else:
    nest.assert_same_structure(regularizers, state_size)
    flat_regularizer = nest.flatten(regularizers)
    if not all([callable(regularizer) for regularizer in flat_regularizer]):
      raise ValueError("Not all the passed regularizers are callable objects.")

  # Produce names for the variables. In the case of a tuple or nested tuple,
  # this is just a sequence of numbers, but for a flat `namedtuple`, we use
  # the field names. NOTE: this could be extended to nested `namedtuple`s,
  # but for now that's extra complexity that's not used anywhere.
  name_prefix = name or "initial_state"
  try:
    name_suffixes = [
        state_size._fields[i] for i in xrange(len(flat_state_size))]
  except (AttributeError, IndexError):
    name_suffixes = range(len(flat_state_size))

  flat_initial_state = []

  for name_suffix, size, init, regularizer in zip(
      name_suffixes, flat_state_size, flat_initializer, flat_regularizer):
    shape_with_batch_dim = [1] + tf.TensorShape(size).as_list()

    variable_name = "{}_{}".format(name_prefix, name_suffix)
    initial_state_module = basic.TrainableVariable(
        shape_with_batch_dim, dtype=dtype, initializers={"w": init},
        regularizers={"w": regularizer}, name=variable_name)
    initial_state_variable = initial_state_module()

    tiled_name = "state_{}_tiled".format(name_suffix)

    initial_state_variable_dims = initial_state_variable.get_shape().ndims
    tile_dims = [batch_size] + [1] * (initial_state_variable_dims - 1)
    flat_initial_state.append(
        tf.tile(initial_state_variable, tile_dims, name=tiled_name))

  return nest.pack_sequence_as(structure=state_size,
                               flat_sequence=flat_initial_state)


Instruction:
Can you write a function in Python 3 where it
returns a decorator to copy documentation from the given function.

Response:
def with_doc(fn_with_doc_to_copy):
  """Returns a decorator to copy documentation from the given function.

  Docstring is copied, including *args and **kwargs documentation.

  Args:
    fn_with_doc_to_copy: Function whose docstring, including *args and
      **kwargs documentation, is to be copied.

  Returns:
    Decorated version of `wrapper_init` with documentation copied from
    `fn_with_doc_to_copy`.
  """

  def decorator(wrapper_init):
    # Wrap the target class's constructor (to assume its docstring),
    # but invoke the wrapper class's constructor.
    @wrapt.decorator
    def wrapping_fn(unused_wrapped, instance, args, kwargs):
      wrapper_init(instance, *args, **kwargs)
    return wrapping_fn(fn_with_doc_to_copy)  # pylint: disable=no-value-for-parameter

  return decorator


Instruction:
Explain what the following Python 3 code does
def _build(self):
    flat_initial_state = nest.flatten(self._initial_state)
    if self._mask is not None:
      flat_mask = nest.flatten(self._mask)
      flat_learnable_state = [
          _single_learnable_state(state, state_id=i, learnable=mask)
          for i, (state, mask) in enumerate(zip(flat_initial_state, flat_mask))]
    else:
      flat_learnable_state = [_single_learnable_state(state, state_id=i)
                              for i, state in enumerate(flat_initial_state)]

    return nest.pack_sequence_as(structure=self._initial_state,
                                 flat_sequence=flat_learnable_state)

Response: Connects the module to the graph.


Instruction:
Can you tell what is the following Python 3 function doing
def _build(self, inputs):

    if self._axis is None:
      axis = list(range(1, inputs.shape.ndims))
    else:
      axis = self._axis

    original_dtype = inputs.dtype
    if original_dtype in [tf.float16, tf.bfloat16]:
      inputs = tf.cast(inputs, tf.float32)

    if inputs.get_shape().ndims < 2:
      raise base.NotSupportedError(
          "Layer normalization expects inputs of at least rank 2."
          " Got inputs of rank {}.".format(inputs.get_shape().ndims))

    # Shape for the learnable scale and offset is the number of channels. See
    # https://arxiv.org/pdf/1803.08494.pdf around equation 6.
    params_shape = inputs.get_shape()[-1:]

    if self._scale:
      if self.GAMMA not in self._initializers:
        self._initializers[self.GAMMA] = create_gamma_initializer()
      self._gamma = tf.get_variable(
          self.GAMMA,
          shape=params_shape,
          dtype=inputs.dtype,
          initializer=self._initializers[self.GAMMA],
          partitioner=self._partitioners.get(self.GAMMA),
          regularizer=self._regularizers.get(self.GAMMA))
    else:
      self._gamma = None

    if self._offset:
      if self.BETA not in self._initializers:
        self._initializers[self.BETA] = create_beta_initializer()
      self._beta = tf.get_variable(
          self.BETA,
          shape=params_shape,
          dtype=inputs.dtype,
          initializer=self._initializers[self.BETA],
          partitioner=self._partitioners.get(self.BETA),
          regularizer=self._regularizers.get(self.BETA))
    else:
      self._beta = None

    mean, var = tf.nn.moments(inputs, axis, keep_dims=True)

    normalized = tf.nn.batch_normalization(inputs, mean, var, self._beta,
                                           self._gamma, self._eps)

    if original_dtype in [tf.float16, tf.bfloat16]:
      normalized = tf.cast(normalized, dtype=original_dtype)
    return normalized

Response: Connects the LayerNorm module into the graph.


Instruction:
Can you tell what is the following Python 3 function doing
def _build(self, memory, query, memory_mask=None):
    if len(memory.get_shape()) != 3:
      raise base.IncompatibleShapeError(
          "memory must have shape [batch_size, memory_size, memory_word_size].")

    if len(query.get_shape()) != 2:
      raise base.IncompatibleShapeError(
          "query must have shape [batch_size, query_word_size].")

    if memory_mask is not None and len(memory_mask.get_shape()) != 2:
      raise base.IncompatibleShapeError(
          "memory_mask must have shape [batch_size, memory_size].")

    # Ensure final dimensions are defined, else the attention logit module will
    # be unable to infer input size when constructing variables.
    inferred_memory_word_size = memory.get_shape()[2].value
    inferred_query_word_size = query.get_shape()[1].value
    if inferred_memory_word_size is None or inferred_query_word_size is None:
      raise base.UnderspecifiedError(
          "memory_word_size and query_word_size must be known at graph "
          "construction time.")

    memory_shape = tf.shape(memory)
    batch_size = memory_shape[0]
    memory_size = memory_shape[1]

    query_shape = tf.shape(query)
    query_batch_size = query_shape[0]

    # Transform query to have same number of words as memory.
    #
    # expanded_query: [batch_size, memory_size, query_word_size].
    expanded_query = tf.tile(tf.expand_dims(query, dim=1), [1, memory_size, 1])

    # Compute attention weights for each memory slot.
    #
    # attention_weight_logits: [batch_size, memory_size]
    with tf.control_dependencies(
        [tf.assert_equal(batch_size, query_batch_size)]):
      concatenated_embeddings = tf.concat(
          values=[memory, expanded_query], axis=2)

    batch_apply_attention_logit = basic.BatchApply(
        self._attention_logit_mod, n_dims=2, name="batch_apply_attention_logit")
    attention_weight_logits = batch_apply_attention_logit(
        concatenated_embeddings)

    # Note: basic.BatchApply() will automatically reshape the [batch_size *
    # memory_size, 1]-shaped result of self._attention_logit_mod(...) into a
    # [batch_size, memory_size, 1]-shaped Tensor. If
    # self._attention_logit_mod(...) returns something with more dimensions,
    # then attention_weight_logits will have extra dimensions, too.
    if len(attention_weight_logits.get_shape()) != 3:
      raise base.IncompatibleShapeError(
          "attention_weight_logits must be a rank-3 Tensor. Are you sure that "
          "attention_logit_mod() returned [batch_size * memory_size, 1]-shaped"
          " Tensor?")

    # Remove final length-1 dimension.
    attention_weight_logits = tf.squeeze(attention_weight_logits, [2])

    # Mask out ignored memory slots by assigning them very small logits. Ensures
    # that every example has at least one valid memory slot, else we'd end up
    # averaging all memory slots equally.
    if memory_mask is not None:
      num_remaining_memory_slots = tf.reduce_sum(
          tf.cast(memory_mask, dtype=tf.int32), axis=[1])
      with tf.control_dependencies(
          [tf.assert_positive(num_remaining_memory_slots)]):
        finfo = np.finfo(np.float32)
        kept_indices = tf.cast(memory_mask, dtype=tf.float32)
        ignored_indices = tf.cast(tf.logical_not(memory_mask), dtype=tf.float32)
        lower_bound = finfo.max * kept_indices + finfo.min * ignored_indices
        attention_weight_logits = tf.minimum(attention_weight_logits,
                                             lower_bound)

    # attended_memory: [batch_size, memory_word_size].
    attention_weight = tf.reshape(
        tf.nn.softmax(attention_weight_logits),
        shape=[batch_size, memory_size, 1])
    # The multiplication is elementwise and relies on broadcasting the weights
    # across memory_word_size. Then we sum across the memory slots.
    attended_memory = tf.reduce_sum(memory * attention_weight, axis=[1])

    # Infer shape of result as much as possible.
    inferred_batch_size, _, inferred_memory_word_size = (
        memory.get_shape().as_list())
    attended_memory.set_shape([inferred_batch_size, inferred_memory_word_size])

    return AttentionOutput(
        read=attended_memory,
        weights=tf.squeeze(attention_weight, [2]),
        weight_logits=attention_weight_logits)

Response: Builds the list of attention outputs for a differentiable read.


Instruction:
Write a Python 3 script for
building the statistics part of the graph when using moving mean and variance.

Response:
def _build_statistics(self, input_batch, use_batch_stats, stat_dtype):
    """Builds the statistics part of the graph when using moving variance.

    Args:
      input_batch: Input batch Tensor.
      use_batch_stats: Boolean to indicate if batch statistics should be
        calculated, otherwise moving averages are returned.
      stat_dtype: TensorFlow datatype to use for the moving mean and variance.

    Returns:
      Tuple of (mean, variance), each of the same datatype as `input_batch`.
    """
    # Set up our moving statistics. When connecting in parallel, this is shared.
    if self.MOVING_MEAN not in self._initializers:
      self._initializers[self.MOVING_MEAN] = create_mean_initializer()
    self._moving_mean = tf.get_variable(
        "moving_mean",
        dtype=stat_dtype,
        shape=(self._num_channels,),
        collections=[
            tf.GraphKeys.MOVING_AVERAGE_VARIABLES,
            tf.GraphKeys.GLOBAL_VARIABLES,
        ],
        initializer=self._initializers[self.MOVING_MEAN],
        trainable=False)

    if self.MOVING_VARIANCE not in self._initializers:
      self._initializers[self.MOVING_VARIANCE] = create_variance_initializer()
    self._moving_variance = tf.get_variable(
        "moving_variance",
        dtype=stat_dtype,
        shape=(self._num_channels,),
        collections=[
            tf.GraphKeys.MOVING_AVERAGE_VARIABLES,
            tf.GraphKeys.GLOBAL_VARIABLES,
        ],
        initializer=self._initializers[self.MOVING_VARIANCE],
        trainable=False)

    def build_batch_stats():
      """Builds the batch statistics calculation ops."""
      mean, variance = tf.nn.moments(input_batch, self._axis,
                                     keep_dims=True, name="normalize_moments")

      return mean, variance

    def build_moving_stats():
      """Retrieves the moving statistics."""
      # If necessary, cast the moving statistics to match the input type.
      # This is required by tf.nn.batch_normalization.
      input_dtype = input_batch.dtype.base_dtype
      if stat_dtype == input_dtype:
        return (
            tf.identity(self._moving_mean),
            tf.identity(self._moving_variance),
        )
      else:
        return (
            tf.cast(self._moving_mean, input_dtype),
            tf.cast(self._moving_variance, input_dtype),
        )

    mean, variance = utils.smart_cond(
        use_batch_stats,
        build_batch_stats,
        build_moving_stats,
    )

    return mean, variance


Instruction:
Here you have a function in Python 3, explain what it does
def _build_update_ops(self, mean, variance, is_training):

    def build_update_ops():
      """Builds the exponential moving average update ops."""

      update_mean_op = moving_averages.assign_moving_average(
          variable=self._moving_mean,
          value=tf.reshape(mean, (self._num_channels,)),
          decay=self._decay_rate,
          zero_debias=False,
          name="update_moving_mean").op

      update_variance_op = moving_averages.assign_moving_average(
          variable=self._moving_variance,
          value=tf.reshape(variance, (self._num_channels,)),
          decay=self._decay_rate,
          zero_debias=False,
          name="update_moving_variance").op

      return update_mean_op, update_variance_op

    def build_no_ops():
      return (tf.no_op(), tf.no_op())

    # Only make the ops if we know that `is_training=True`, or the value of
    # `is_training` is unknown.
    is_training_const = utils.constant_value(is_training)
    if is_training_const is None or is_training_const:
      update_mean_op, update_variance_op = utils.smart_cond(
          is_training,
          build_update_ops,
          build_no_ops,
      )
      return (update_mean_op, update_variance_op)
    else:
      return None

Response: Builds the update ops for moving average.


Instruction:
Create a Python 3 function to
create a fused batch normalization op.

Response:
def _fused_batch_norm_op(self, input_batch, mean, variance, use_batch_stats):
    """Creates a fused batch normalization op."""
    # Store the original shape of the mean and variance.
    mean_shape = mean.get_shape()
    variance_shape = variance.get_shape()
    # The fused batch norm expects the mean, variance, gamma and beta
    # tensors to have dimension 1, so we flatten them to remove the
    # extra dimensions. In addition, it expects the input_batch to have
    # dimension 4, so we reshape it accordingly.
    gamma_flatten = tf.reshape(self._gamma, shape=(self._num_channels,))
    beta_flatten = tf.reshape(self._beta, shape=(self._num_channels,))
    flatten_mean = tf.reshape(mean, shape=(self._num_channels,))
    flatten_variance = tf.reshape(variance, shape=(self._num_channels,))
    use_batch_stats = tf.convert_to_tensor(use_batch_stats)

    input_shape = input_batch.get_shape()
    output_shape = [-1] + input_shape.as_list()[1:]

    flat_image_size = np.prod(self._image_shape, dtype=np.int32)

    if len(self._data_format) == 4:
      fusable_data_format = self._data_format
      fusable_batch = input_batch
    elif self._channel_index == 1 and self._image_shape:
      fusable_data_format = "NCHW"
      fusable_batch = tf.reshape(
          input_batch,
          shape=(-1, self._num_channels, 1, flat_image_size))
    else:
      # The CPU implementation of FusedBatchNorm only supports NHWC tensor
      # format for now.
      fusable_data_format = "NHWC"
      fusable_batch = tf.reshape(
          input_batch,
          shape=(-1, 1, flat_image_size, self._num_channels))

    common_args = {
        "scale": gamma_flatten,
        "offset": beta_flatten,
        "epsilon": self._eps,
        "data_format": fusable_data_format,
        "name": "batch_norm"
    }

    def use_batch_stats_fused_batch_norm():
      return tf.nn.fused_batch_norm(
          fusable_batch,
          mean=None,
          variance=None,
          is_training=True,
          **common_args)

    def moving_average_fused_batch_norm():
      return tf.nn.fused_batch_norm(
          fusable_batch,
          mean=flatten_mean,
          variance=flatten_variance,
          is_training=False,
          **common_args)

    batch_norm_op, mean, variance = utils.smart_cond(
        use_batch_stats, use_batch_stats_fused_batch_norm,
        moving_average_fused_batch_norm)

    if len(self._data_format) != 4:
      batch_norm_op = tf.reshape(batch_norm_op, output_shape)
    mean = tf.reshape(mean, mean_shape)
    variance = tf.reshape(variance, variance_shape)
    return batch_norm_op, mean, variance


Instruction:
Can you generate the documentation for the following Python 3 function
def _build(self,
             input_batch,
             is_training,
             test_local_stats=False):
    """Connects the BatchNormV2 module into the graph.

    Args:
      input_batch: A Tensor of the same dimension as `len(data_format)`.
      is_training: A boolean to indicate if the module should be connected in
        training mode, meaning the moving averages are updated. Can be a Tensor.
      test_local_stats: A boolean to indicate if local batch statistics should
        be used when `is_training=False`. If not, moving averages are used.
        By default `False`. Can be a Tensor.

    Returns:
      A tensor with the same shape as `input_batch`.

    Raises:
      base.IncompatibleShapeError: If `data_format` is not valid for the
        input shape.
      base.NotSupportedError: If `input_batch` has data type of `tf.bfloat16`.
    """
    input_shape = input_batch.get_shape()

    if not self._data_format:
      if len(input_shape) == 2:
        self._data_format = "NC"
      elif len(input_shape) == 3:
        self._data_format = "NWC"
      elif len(input_shape) == 4:
        self._data_format = "NHWC"
      elif len(input_shape) == 5:
        self._data_format = "NDHWC"
      else:
        raise base.IncompatibleShapeError(
            "Input shape {} has too many or too few dimensions.".format(
                input_shape))

    self._channel_index = self._data_format.index("C")
    # Use list to turn range into iterator in python3.
    self._axis = list(range(len(self._data_format)))
    del self._axis[self._channel_index]

    if len(self._data_format) != len(input_shape):
      raise base.IncompatibleShapeError(
          "Incorrect data format {} for input shape {}.".format(
              self._data_format, input_shape))

    dtype = input_batch.dtype.base_dtype
    if self._fused and dtype == tf.bfloat16:
      raise base.NotSupportedError(
          "Fused batch norm does not support tf.bfloat16.")
    # Maintain moving averages at a minimum precision of tf.float32.
    stat_dtype = tf.float32 if dtype in [tf.float16, tf.bfloat16] else dtype

    self._num_channels = int(input_shape[self._channel_index])
    if self._channel_index == 1:
      self._image_shape = [int(x) for x in input_shape[2:]]
    else:
      self._image_shape = [int(x) for x in input_shape[1:-1]]

    self._expanded_mean_shape = [1] * len(input_shape)
    self._expanded_mean_shape[self._channel_index] = self._num_channels

    use_batch_stats = is_training | test_local_stats

    mean, variance = self._build_statistics(input_batch, use_batch_stats,
                                            stat_dtype)

    # Sets up optional gamma and beta parameters
    self._build_scale_offset(dtype)
    # Sets up the batch normalization op.
    out, mean, variance = self._batch_norm_op(input_batch, mean, variance,
                                              use_batch_stats, stat_dtype)
    # Sets up the update op.
    update_ops = self._build_update_ops(mean, variance, is_training)

    # Put update ops in the update ops collection if given, otherwise add as
    # control dependencies of the output.
    if update_ops:
      if self._update_ops_collection:
        for update_op in update_ops:
          tf.add_to_collection(self._update_ops_collection, update_op)
      else:
        with tf.control_dependencies(update_ops):
          out = tf.identity(out)

    return out

Response: Connects the BatchNormV2 module into the graph.


Instruction:
Here you have a function in Python 3, explain what it does
def _range_along_dimension(range_dim, shape):
  rank = len(shape)
  if range_dim >= rank:
    raise ValueError("Cannot calculate range along non-existent index.")
  indices = tf.range(start=0, limit=shape[range_dim])
  indices = tf.reshape(
      indices,
      shape=[1 if i != range_dim else shape[range_dim] for i in range(rank)])
  return tf.tile(indices,
                 [shape[i] if i != range_dim else 1 for i in range(rank)])

Response: Construct a Tensor that counts the distance along a single dimension. This is a helper function for calculating the range along a single dimension.


Instruction:
Can you implement a function in Python 3 that
constructs an initializer function for constructing a noisy identity convolution kernels.

Response:
def noisy_identity_kernel_initializer(base_num_channels, stddev=1e-8):
  """Build an initializer for constructing near-identity convolution kernels.

  Construct a convolution kernel where in_channels and out_channels are
  multiples of base_num_channels, but need not be equal. This initializer is
  essentially the same as identity_kernel_initializer, except that magnitude
  is "spread out" across multiple copies of the input.

  Args:
    base_num_channels: int. Number that divides both in_channels and
      out_channels.
    stddev: float. Standard deviation of truncated normal noise added to
      off-entries to break ties.

  Returns:
    Initializer function for building a noisy identity kernel.
  """

  # pylint: disable=unused-argument
  def _noisy_identity_kernel_initializer(shape,
                                         dtype=tf.float32,
                                         partition_info=None):
    """Constructs a noisy identity kernel.

    Args:
      shape: List of integers. Represents shape of result.
      dtype: data type for values in result.
      partition_info: Partition information for initializer functions. Ignored.

    Returns:
      Tensor of desired shape and dtype such that applying it as a convolution
        kernel results in a noisy near-identity operation.

    Raises:
      ValueError: If shape does not define a valid kernel.
                  If filter width and height differ.
                  If filter width and height are not odd numbers.
                  If number of input and output channels are not multiples of
                    base_num_channels.
    """
    if len(shape) != 4:
      raise ValueError("Convolution kernels must be rank 4.")

    filter_height, filter_width, in_channels, out_channels = shape

    if filter_width != filter_height:
      raise ValueError(
          "Noisy identity initializer only works for square filters.")
    if filter_width % 2 != 1:
      raise ValueError(
          "Noisy identity initializer requires filters have odd height and "
          "width.")
    if (in_channels % base_num_channels != 0 or
        out_channels % base_num_channels != 0):
      raise ValueError("in_channels and out_channels must both be multiples of "
                       "base_num_channels.")

    middle_pixel = filter_height // 2
    is_middle_pixel = tf.logical_and(
        tf.equal(_range_along_dimension(0, shape), middle_pixel),
        tf.equal(_range_along_dimension(1, shape), middle_pixel))
    is_same_channel_multiple = tf.equal(
        tf.floordiv(
            _range_along_dimension(2, shape) * base_num_channels, in_channels),
        tf.floordiv(
            _range_along_dimension(3, shape) * base_num_channels, out_channels))
    noise = tf.truncated_normal(shape, stddev=stddev, dtype=dtype)
    return tf.where(
        tf.logical_and(is_same_channel_multiple, is_middle_pixel),
        tf.ones(
            shape, dtype=dtype) * (base_num_channels / out_channels),
        noise)

  return _noisy_identity_kernel_initializer


Instruction:
Write a Python 3 function that can
build the dilation module.

Response:
def _build(self, images):
    """Build dilation module.

    Args:
      images: Tensor of shape [batch_size, height, width, depth]
        and dtype float32. Represents a set of images with an arbitrary depth.
        Note that when using the default initializer, depth must equal
        num_output_classes.

    Returns:
      Tensor of shape [batch_size, height, width, num_output_classes] and dtype
        float32. Represents, for each image and pixel, logits for per-class
        predictions.

    Raises:
      IncompatibleShapeError: If images is not rank 4.
      ValueError: If model_size is not one of 'basic' or 'large'.
    """
    num_classes = self._num_output_classes

    if len(images.get_shape()) != 4:
      raise base.IncompatibleShapeError(
          "'images' must have shape [batch_size, height, width, depth].")

    if self.WEIGHTS not in self._initializers:
      if self._model_size == self.BASIC:
        self._initializers[self.WEIGHTS] = identity_kernel_initializer
      elif self._model_size == self.LARGE:
        self._initializers[self.WEIGHTS] = noisy_identity_kernel_initializer(
            num_classes)
      else:
        raise ValueError("Unrecognized model_size: %s" % self._model_size)

    if self.BIASES not in self._initializers:
      self._initializers[self.BIASES] = tf.zeros_initializer()

    if self._model_size == self.BASIC:
      self._conv_modules = [
          self._dilated_conv_layer(num_classes, 1, True, "conv1"),
          self._dilated_conv_layer(num_classes, 1, True, "conv2"),
          self._dilated_conv_layer(num_classes, 2, True, "conv3"),
          self._dilated_conv_layer(num_classes, 4, True, "conv4"),
          self._dilated_conv_layer(num_classes, 8, True, "conv5"),
          self._dilated_conv_layer(num_classes, 16, True, "conv6"),
          self._dilated_conv_layer(num_classes, 1, True, "conv7"),
          self._dilated_conv_layer(num_classes, 1, False, "conv8"),
      ]
    elif self._model_size == self.LARGE:
      self._conv_modules = [
          self._dilated_conv_layer(2 * num_classes, 1, True, "conv1"),
          self._dilated_conv_layer(2 * num_classes, 1, True, "conv2"),
          self._dilated_conv_layer(4 * num_classes, 2, True, "conv3"),
          self._dilated_conv_layer(8 * num_classes, 4, True, "conv4"),
          self._dilated_conv_layer(16 * num_classes, 8, True, "conv5"),
          self._dilated_conv_layer(32 * num_classes, 16, True, "conv6"),
          self._dilated_conv_layer(32 * num_classes, 1, True, "conv7"),
          self._dilated_conv_layer(num_classes, 1, False, "conv8"),
      ]
    else:
      raise ValueError("Unrecognized model_size: %s" % self._model_size)

    dilation_mod = sequential.Sequential(self._conv_modules, name="dilation")
    return dilation_mod(images)


Instruction:
Can you implement a function in Python 3 that
creates a dilated convolution layer.

Response:
def _dilated_conv_layer(self, output_channels, dilation_rate, apply_relu,
                          name):
    """Create a dilated convolution layer.

    Args:
      output_channels: int. Number of output channels for each pixel.
      dilation_rate: int. Represents how many pixels each stride offset will
        move. A value of 1 indicates a standard convolution.
      apply_relu: bool. If True, a ReLU non-linearlity is added.
      name: string. Name for layer.

    Returns:
      a sonnet Module for a dilated convolution.
    """
    layer_components = [
        conv.Conv2D(
            output_channels, [3, 3],
            initializers=self._initializers,
            regularizers=self._regularizers,
            rate=dilation_rate,
            name="dilated_conv_" + name),
    ]
    if apply_relu:
      layer_components.append(lambda net: tf.nn.relu(net, name="relu_" + name))
    return sequential.Sequential(layer_components, name=name)


Instruction:
Write a Python 3 script for
wrapping the function and prints a warning - once per extra_message.

Response:
def with_deprecation_warning(fn, extra_message=''):
  """Wraps the function and prints a warn-once (per `extra_message`) warning."""
  def new_fn(*args, **kwargs):
    if extra_message not in _DONE_WARN:
      tf.logging.warning(
          'Sonnet nest is deprecated. Please use '
          'tf.contrib.framework.nest instead. '
          + extra_message
      )
      _DONE_WARN[extra_message] = True
    return fn(*args, **kwargs)
  return new_fn


Instruction:
Can you tell what is the following Python 3 function doing
def custom_build(inputs, is_training, keep_prob):
  outputs = snt.Conv2D(output_channels=32, kernel_shape=4, stride=2)(inputs)
  outputs = snt.BatchNorm()(outputs, is_training=is_training)
  outputs = tf.nn.relu(outputs)
  outputs = snt.Conv2D(output_channels=64, kernel_shape=4, stride=2)(outputs)
  outputs = snt.BatchNorm()(outputs, is_training=is_training)
  outputs = tf.nn.relu(outputs)
  outputs = snt.BatchFlatten()(outputs)
  outputs = tf.nn.dropout(outputs, keep_prob=keep_prob)
  outputs = snt.Linear(output_size=10)(outputs)
  return outputs

Response: A custom build method to wrap into a sonnet Module.


Instruction:
Make a summary of the following Python 3 code
def _get_or_create_stack(name):
  stack = getattr(_LOCAL_STACKS, name, None)
  if stack is None:
    stack = []
    setattr(_LOCAL_STACKS, name, stack)
  return stack

Response: Returns a thread local stack uniquified by the given name.


Instruction:
Given the following Python 3 function, write the documentation
def bayes_by_backprop_getter(
    posterior_builder=diagonal_gaussian_posterior_builder,
    prior_builder=fixed_gaussian_prior_builder,
    kl_builder=stochastic_kl_builder,
    sampling_mode_tensor=None,
    fresh_noise_per_connection=True,
    keep_control_dependencies=False):
  """Creates a custom getter which does Bayes by Backprop.

  Please see `tf.get_variable` for general documentation on custom getters.

  All arguments are optional. If nothing is configued, then a diagonal gaussian
  posterior will be used, and a fixed N(0, 0.01) prior will be used. Please
  see the default `posterior_builder` and `prior_builder` for a more detailed
  understanding of the default settings.

  Args:
    posterior_builder: A builder function which constructs an instance of
      `tfp.distributions.Distribution` which shall serve as the posterior over
      the `tf.Variable` of interest. The builder receives the `getter` and the
      arguments forwarded from `tf.get_variable`. Suppose one wrote

      ```
      tf.get_variable(
          'weights', shape=(3,), initializer=tf.zeros_initializer,
          dtype=tf.float32)
      ```

      then the `posterior_builder` argument would receive the `name`, `shape`,
      `initializer`, and `dtype` arguments passed above. The builder must return
      a `tfp.distributions.Distribution` object.

      Please see the `tf.get_variable` for documentation on `custom_getter` and
      `getter`, and see `bbb.diagonal_gaussian_posterior_builder`
      (the default) for an example of using this builder API.
    prior_builder: A builder function which constructs an instance of
      `tfp.distributions.Distribution` which shall serve as the prior over the
      `tf.Variable` of interest. Identical API to `posterior_builder`. See
      `bbb.fixed_gaussian_prior_builder` (the default) for an example.
    kl_builder: A builder function which receives the posterior distribution,
      prior distribution, and a sample from the posterior. It returns a
      scalar-shaped `tf.Tensor` representing the total KL cost for the
      `tf.Variable` in question. See `bbb.stochastic_kl_builder` (default) and
      `bbb.analytic_kl_builder` for examples.
    sampling_mode_tensor: A `tf.Tensor` which determines how an estimate from
      the posterior is produced. It must be scalar-shaped and have a `dtype` of
      `tf.string`. Valid values for this tensor are `bbb.EstimatorModes.sample`
      (which is the default), `bbb.EstimatorModes.mean`, and
      `bbb.EstimatorModes.last_sample`. `bbb.EstimatorModes.sample` is
      appropriate for training, and `bbb.EstimatorModes.mean` can be used
      at test time.
    fresh_noise_per_connection: A boolean. Indicates that each time a stochastic
      variable is retrieved with this custom getter, new sampling noise should
      be used. This is `True` by default. If this argument is set to `False`,
      then the same noise is used for each connection. Note that this does not
      apply to connections within a `tf.while_loop`; the same sampling noise
      is always used in different iterations of a `tf.while_loop` within one
      `session.run()` call. See the unit tests for details.
    keep_control_dependencies: A boolean. This argument should only be
      used by advanced users. Indicates that each time a stochastic variable is
      retrieved in the loop body of a `tf.while_loop` construct, new sampling
      noise should be used.
      The default behavior is `False`, so that RNNs use the same weights at each
      recurrent time step. This is done by removing the creation of the Variable
      from any existing control flow contexts. Notably, the Variables will be
      created outside the context of any tf.while_loop, making them fetchable.
      When this argument is `True`, any Variables used in the loop body of a
      `tf.while_loop` will be non-fetchable. If the KL cost needs to be
      evaluated, the Variable must *first* be used *outside* the loop body. This
      op using the Variable simply needs to be placed on the graph to get a
      stochastic estimate of the KL; it doesn't need to ever be used. Example:

      ```
      def loop_body(i):
        logits = sonnet_module(queue)
        i = i + 1

      with tf.variable_scope('bbb', custom_getter=bbb.bayes_by_backprop_getter(
          fresh_noise_per_connection=True,
          keep_control_dependencies=True)):
        unused_op = sonnet_module(queue)  # Adds KL estimate to bbb Collection
        final_i = tf.while_loop(lambda i: i < 5, loop_body, tf.constant(0.))
      ```

      Here when we add `unused_op` to the graph, we also add a number of tensors
      associated with the particular stochastic variable, including its
      contribution to the KL cost, to a graph-level registry. These are
      organized in a per-stochastic-variable data structure and be accessed with
      `bbb.get_variable_metadata()`. Without this line, these Tensors would
      instead be added the first time the Variable is used in the while_loop,
      which would make them non-fetchable.

      In all cases, the KL cost is only added once per Variable, which is the
      correct behavior, since if a variable is used multiple times in a model,
      the KL cost should remain unaffected.

  Returns:
    A `custom_getter` function which implements Bayes by Backprop.
  """

  if sampling_mode_tensor is None:
    sampling_mode_tensor = tf.constant(EstimatorModes.sample)

  def custom_getter(getter, name, *args, **kwargs):
    """The custom getter that will be returned."""
    if not kwargs.get("trainable", True):
      return getter(name, *args, **kwargs)
    if kwargs["dtype"] not in _OK_DTYPES_FOR_BBB:
      raise ValueError("Disallowed data type {}.".format(kwargs["dtype"]))

    var_scope = tf.get_variable_scope()
    if var_scope.reuse and not fresh_noise_per_connection:
      # Re-use the sampling noise by returning the very same posterior sample
      # if configured to do so.
      the_match = [
          x for x in get_variable_metadata() if x.raw_variable_name == name]
      if not the_match:
        raise ValueError(
            "Internal error. No metadata for variable {}".format(name))
      if len(the_match) > 1:
        raise ValueError(
            "Multiple matches for variable {}. Matches: {}".format(
                name, [x.raw_variable_name for x in the_match]))

      return the_match[0].posterior_estimate

    raw_variable_shape = kwargs["shape"]

    def construct_subgraph():
      """Constructs subgraph used to reparameterize the variable in question."""
      posterior = posterior_builder(
          getter,
          name=name,
          *args, **kwargs)
      prior = prior_builder(
          getter,
          name=name,
          *args, **kwargs)

      # If the user does not return an extra dictionary of prior variables,
      # then fill in an empty dictionary.
      if isinstance(posterior, collections.Sequence):
        posterior_dist, posterior_vars = posterior
      else:
        posterior_dist, posterior_vars = posterior, {}

      if isinstance(prior, collections.Sequence):
        prior_dist, prior_vars = prior
      else:
        prior_dist, prior_vars = prior, {}

      if posterior_dist.reparameterization_type != _OK_PZATION_TYPE:
        raise ValueError(
            "Distribution {} incompatible with Bayes by Backprop.".format(
                posterior_dist.__class__.__name__))

      posterior_estimator = _produce_posterior_estimate(posterior_dist,
                                                        sampling_mode_tensor,
                                                        name)
      kl_cost = kl_builder(posterior_dist, prior_dist, posterior_estimator)
      variable_metadata = _VariableMetadata(
          raw_variable_name=name,
          raw_variable_shape=raw_variable_shape,
          scope_name=var_scope.name,
          posterior=posterior_dist,
          posterior_estimate=posterior_estimator,
          prior=prior_dist,
          kl_cost=kl_cost,
          prior_vars=prior_vars,
          posterior_vars=posterior_vars)
      return posterior_estimator, variable_metadata

    # Entering the `tf.control_dependencies(None)` context is crucial to
    # provide compatibility with `tf.while_loop` and thus RNNs. The main thing
    # it does is making the `kl_cost` fetchable by causing these ops to be
    # created outside the context of any tf.while_loop. Note also that it causes
    # a RNN core's weights to be sampled just once when unrolled over a
    # sequence, rather than at every timestep.
    control_deps = [] if keep_control_dependencies else None
    with tf.control_dependencies(control_deps):
      posterior_estimator, var_metadata = construct_subgraph()

    # Only add these ops to a collection once per unique variable.
    # This is to ensure that KL costs are not tallied up more than once.
    var_with_name = _all_var_metadata_registry[tf.get_default_graph()].get(name)
    if var_with_name is None:
      _all_var_metadata_registry[tf.get_default_graph()][name] = var_metadata

    return posterior_estimator
  return custom_getter

Response: Creates a custom getter which does Bayes by Backprop.


Instruction:
Can you tell what is the following Python 3 function doing
def _produce_posterior_estimate(posterior_dist, posterior_estimate_mode,
                                raw_var_name):
  """Create tensor representing estimate of posterior.

  Args:
    posterior_dist: An instance of `tfp.distributions.Distribution`.
        The variational posterior from which to produce an estimate of the
        variable in question.
    posterior_estimate_mode: A `Tensor` of dtype `tf.string`, which
        determines the inference mode.
    raw_var_name: The name of the variable over which inference is done.

  Returns:
    `z_sample`, a `Tensor` representing an estimate derived from the
        posterior distribution.
  """
  conds = [
      tf.equal(posterior_estimate_mode,
               tf.constant(EstimatorModes.sample),
               name="equal_sample_mode"),
      tf.equal(posterior_estimate_mode,
               tf.constant(EstimatorModes.mean),
               name="equal_mean_mode"),
      tf.equal(posterior_estimate_mode,
               tf.constant(EstimatorModes.last_sample),
               name="equal_last_sample_mode"),
  ]
  # pylint: disable=unnecessary-lambda
  results = [
      lambda: posterior_dist.sample(),
      lambda: posterior_dist.mean(),
      lambda: posterior_dist.last_sample()
  ]

  def default_case_branch_raising_error():
    err_msg = "Invalid posterior estimate mode."
    raise_err = tf.Assert(tf.constant(False), data=[tf.constant(err_msg)])
    with tf.control_dependencies([raise_err]):
      return posterior_dist.mean()

  if hasattr(posterior_dist, "last_sample"):
    cases = {conds[0]: results[0], conds[1]: results[1], conds[2]: results[2]}
  else:
    cases = {conds[0]: results[0], conds[1]: results[1]}
  z_sample = tf.case(
      cases,
      exclusive=True,
      default=default_case_branch_raising_error,
      name="{}_posterior_estimate".format(raw_var_name))
  # pylint: enable=unnecessary-lambda
  return z_sample

Response: Create a tensor representing an estimate of the posterior.


Instruction:
Can you generate the documentation for the following Python 3 function
def get_total_kl_cost(name="total_kl_cost", filter_by_name_substring=None):
  all_variable_metadata = get_variable_metadata(filter_by_name_substring)
  if not all_variable_metadata:
    tf.logging.warning("No Bayes by Backprop variables found!")
    return tf.constant(0.0, shape=())
  return tf.add_n([md.kl_cost for md in all_variable_metadata], name=name)

Response: Returns the total KL cost for all variables in the current graph.


Instruction:
Explain what the following Python 3 code does
def output_shape(self):
    return (self._block_shape[0] * self._block_rows,
            self._block_shape[1] * self._block_rows)

Response: The shape of the output matrix.


Instruction:
Here you have a function in Python 3, explain what it does
def _left_zero_blocks(self, r):
    if not self._include_off_diagonal:
      return r
    elif not self._upper:
      return 0
    elif self._include_diagonal:
      return r
    else:
      return r + 1

Response: Number of blocks with zeros from the left in block row r.


Instruction:
Given the following Python 3 function, write the documentation
def _right_zero_blocks(self, r):
    if not self._include_off_diagonal:
      return self._block_rows - r - 1
    elif self._upper:
      return 0
    elif self._include_diagonal:
      return self._block_rows - r - 1
    else:
      return self._block_rows - r

Response: Number of blocks with zeros from the right in block row r.


Instruction:
Can you implement a function in Python 3 that
numbers of content blocks in block row r.

Response:
def _content_blocks(self, r):
    """Number of content blocks in block row `r`."""
    return (self._block_rows - self._left_zero_blocks(r)
            - self._right_zero_blocks(r))


Instruction:
Write a Python 3 function that can
construct the data model loss and optimizer then train.

Response:
def build_and_train(iterations, log_stride, test=False):
  """Construct the data, model, loss and optimizer then train."""

  # Test mode settings.
  batch_size = 2 if test else FLAGS.batch_size
  num_mems = 2 if test else FLAGS.num_mems
  num_heads = 1 if test else FLAGS.num_mems
  num_blocks = 1 if test else FLAGS.num_mems
  head_size = 4 if test else FLAGS.head_size
  num_objects = 2 if test else FLAGS.num_objects
  num_features = 4 if test else FLAGS.num_features
  mlp_size = (20,) if test else (256, 256, 256, 256)

  with tf.Graph().as_default():
    t0 = time.time()

    # Initialize the dataset.
    dataset = dataset_nth_farthest.NthFarthest(
        batch_size, num_objects, num_features)

    # Create the model.
    core = snt.RelationalMemory(
        mem_slots=num_mems,
        head_size=head_size,
        num_heads=num_heads,
        num_blocks=num_blocks,
        gate_style=FLAGS.gate_style)

    final_mlp = snt.nets.MLP(
        output_sizes=mlp_size,
        activate_final=True)

    model = SequenceModel(
        core=core,
        target_size=num_objects,
        final_mlp=final_mlp)

    tf.logging.info("Instantiated models ({:3f})".format(time.time() - t0))

    # Get train and test data.
    inputs_train, labels_train = dataset.get_batch()
    inputs_test, labels_test = dataset.get_batch()

    # Define target accuracy.
    def compute_accuracy(logits, targets, name="accuracy"):
      correct_pred = tf.cast(
          tf.equal(tf.cast(targets, tf.int64), tf.argmax(logits, 1)),
          tf.float32)
      return tf.reduce_mean(correct_pred, name=name)

    # Define the loss & accuracy.
    def loss_fn(inputs, labels):
      """Creates the loss and the exports."""
      logits = model(inputs)
      labels = tf.cast(labels, tf.int32)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=logits, labels=labels))
      accuracy = compute_accuracy(logits, labels)
      return loss, accuracy

    # Get training step counter.
    global_step = tf.get_variable(
        name="global_step",
        shape=[],
        dtype=tf.int64,
        initializer=tf.zeros_initializer(),
        trainable=False,
        collections=[tf.GraphKeys.GLOBAL_VARIABLES, tf.GraphKeys.GLOBAL_STEP])

    # Create the optimizer.
    learning_rate_op = tf.reduce_max([
        tf.train.exponential_decay(
            FLAGS.learning_rate,
            global_step,
            decay_steps=FLAGS.epochs // 100,
            decay_rate=0.9,
            staircase=False),
        FLAGS.min_learning_rate
    ])
    optimizer = tf.train.AdamOptimizer(learning_rate_op)
    train_loss, _ = loss_fn(inputs_train, labels_train)
    step_op = optimizer.minimize(train_loss, global_step=global_step)

    # Compute test accuracy
    logits_test = model(inputs_test)
    labels_test = tf.cast(labels_test, tf.int32)
    test_acc = compute_accuracy(logits_test, labels_test)

    tf.logging.info("Created losses and optimizers ({:3f})".format(
        time.time() - t0))

    # Begin Training.
    t0 = time.time()
    train_losses = []
    steps = []
    test_accs = []
    tf.logging.info("Starting training ({:3f})".format(time.time() - t0))
    with tf.train.SingularMonitoredSession() as sess:
      for it in six.moves.range(iterations):
        sess.run([step_op, learning_rate_op])
        if it % log_stride == 0:
          loss_v, acc_v = sess.run([train_loss, test_acc])
          elapsed = time.time() - t0
          tf.logging.info(
              "iter: {:2d}, train loss {:3f}; test acc {:3f} ({:3f})".format(
                  it, loss_v, acc_v, elapsed))
          train_losses.append(loss_v)
          steps.append(it)
        test_accs.append(acc_v)
  return steps, train_losses, test_accs


Instruction:
Can you generate a brief explanation for the following Python 3 code
def _clip_gradient_op(dtype):

  def clip_gradient_backward(op, grad):
    clip_value_min = op.inputs[1]
    clip_value_max = op.inputs[2]
    clipped_grad = tf.clip_by_value(grad, clip_value_min, clip_value_max)
    return clipped_grad, None, None

  def clip_gradient_forward(x, clip_value_min, clip_value_max):
    del clip_value_min  # Unused.
    del clip_value_max  # Unused.
    return x

  func_name = "ClipGradient_{}".format(dtype.name)
  return function.Defun(
      dtype, dtype, dtype,
      python_grad_func=clip_gradient_backward,
      func_name=func_name)(clip_gradient_forward)

Response: Create an op that clips gradients using a Defun.


Instruction:
How would you explain what the following Python 3 function does
def clip_gradient(net, clip_value_min, clip_value_max, name=None):
  if not net.dtype.is_floating:
    raise ValueError("clip_gradient does not support non-float `net` inputs.")

  with tf.name_scope(name, "clip_gradient", values=[net]):
    dtype = net.dtype.base_dtype  # Convert ref dtypes to regular dtypes.
    min_tensor = tf.convert_to_tensor(clip_value_min, dtype=dtype)
    max_tensor = tf.convert_to_tensor(clip_value_max, dtype=dtype)

    clip_gradient_op = _clip_gradient_op(dtype)
    output = clip_gradient_op(net, min_tensor, max_tensor)
    output.set_shape(net.get_shape())

  return output

Response: Clips respective gradients of a given tensor.


Instruction:
Explain what the following Python 3 code does
def ptb_raw_data(data_path):

  train_path = os.path.join(data_path, "ptb.train.txt")
  valid_path = os.path.join(data_path, "ptb.valid.txt")
  test_path = os.path.join(data_path, "ptb.test.txt")

  word_to_id = _build_vocab(train_path)
  train_data = _file_to_word_ids(train_path, word_to_id)
  valid_data = _file_to_word_ids(valid_path, word_to_id)
  test_data = _file_to_word_ids(test_path, word_to_id)
  return train_data, valid_data, test_data, word_to_id

Response: Loads PTB raw data from data directory data_path.


Instruction:
Can you generate the documentation for the following Python 3 function
def _check_and_assign_normalization_members(self, normalization_ctor,
                                              normalization_kwargs):
    """Checks that the normalization constructor is callable."""
    if isinstance(normalization_ctor, six.string_types):
      normalization_ctor = util.parse_string_to_constructor(normalization_ctor)
    if normalization_ctor is not None and not callable(normalization_ctor):
      raise ValueError(
          "normalization_ctor must be a callable or a string that specifies "
          "a callable, got {}.".format(normalization_ctor))
    self._normalization_ctor = normalization_ctor
    self._normalization_kwargs = normalization_kwargs

Response: Checks that the normalization constructor is callable and assigns the normalization members to the object.


Instruction:
Explain what the following Python 3 code does
def _parse_normalization_kwargs(self, use_batch_norm, batch_norm_config,
                                  normalization_ctor, normalization_kwargs):
    """Sets up normalization, checking old and new flags."""
    if use_batch_norm is not None:
      # Delete this whole block when deprecation is done.
      util.deprecation_warning(
          "`use_batch_norm` kwarg is deprecated. Change your code to instead "
          "specify `normalization_ctor` and `normalization_kwargs`.")
      if not use_batch_norm:
        # Explicitly set to False - normalization_{ctor,kwargs} has precedence.
        self._check_and_assign_normalization_members(normalization_ctor,
                                                     normalization_kwargs or {})
      else:  # Explicitly set to true - new kwargs must not be used.
        if normalization_ctor is not None or normalization_kwargs is not None:
          raise ValueError(
              "if use_batch_norm is specified, normalization_ctor and "
              "normalization_kwargs must not be.")
        self._check_and_assign_normalization_members(batch_norm.BatchNorm,
                                                     batch_norm_config or {})
    else:
      # Old kwargs not set, this block will remain after removing old kwarg.
      self._check_and_assign_normalization_members(normalization_ctor,
                                                   normalization_kwargs or {})

Response: Sets up normalization checking old and new flags.


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

    # Here we are entering the module's variable scope to name our submodules
    # correctly (not to create variables). As such it's safe to not check
    # whether we're in the same graph. This is important if we're constructing
    # the module in one graph and connecting it in another (e.g. with `defun`
    # the module is created in some default graph, and connected to a capturing
    # graph in order to turn it into a graph function).
    with self._enter_variable_scope(check_same_graph=False):
      self._layers = tuple(conv.Conv2D(name="conv_2d_{}".format(i),  # pylint: disable=g-complex-comprehension
                                       output_channels=self._output_channels[i],
                                       kernel_shape=self._kernel_shapes[i],
                                       stride=self._strides[i],
                                       rate=self._rates[i],
                                       padding=self._paddings[i],
                                       use_bias=self._use_bias[i],
                                       initializers=self._initializers,
                                       partitioners=self._partitioners,
                                       regularizers=self._regularizers,
                                       data_format=self._data_format)
                           for i in xrange(self._num_layers))

Response: Instantiates all the convolutional modules used in the network.


Instruction:
Can you generate the documentation for the following Python 3 function
def _build(self, inputs, **normalization_build_kwargs):
    if (self._normalization_ctor in {batch_norm.BatchNorm,
                                     batch_norm_v2.BatchNormV2} and
        "is_training" not in normalization_build_kwargs):
      raise ValueError("Boolean is_training flag must be explicitly specified "
                       "when using batch normalization.")

    self._input_shape = tuple(inputs.get_shape().as_list())
    net = inputs

    final_index = len(self._layers) - 1
    for i, layer in enumerate(self._layers):
      net = layer(net)

      if i != final_index or self._normalize_final:
        if self._normalization_ctor is not None:
          # The name 'batch_norm' is used even if something else like
          # LayerNorm is being used. This is to avoid breaking old checkpoints.
          normalizer = self._normalization_ctor(
              name="batch_norm_{}".format(i),
              **self._normalization_kwargs)

          net = normalizer(
              net, **util.remove_unsupported_kwargs(
                  normalizer, normalization_build_kwargs))
        else:
          if normalization_build_kwargs:
            tf.logging.warning(
                "No normalization configured, but extra kwargs "
                "provided: {}".format(normalization_build_kwargs))

      if i != final_index or self._activate_final:
        net = self._activation(net)

    return net

Response: Assembles the ConvNet2D and connects it to the graph.


Instruction:
Can you tell what is the following Python 3 function doing
def _transpose(self,
                 transpose_constructor,
                 name=None,
                 output_channels=None,
                 kernel_shapes=None,
                 strides=None,
                 paddings=None,
                 activation=None,
                 activate_final=None,
                 normalization_ctor=None,
                 normalization_kwargs=None,
                 normalize_final=None,
                 initializers=None,
                 partitioners=None,
                 regularizers=None,
                 use_bias=None,
                 data_format=None):
    """Returns transposed version of this network.

    Args:
      transpose_constructor: A method that creates an instance of the transposed
        network type. The method must accept the same kwargs as this methods
        with the exception of the `transpose_constructor` argument.
      name: Optional string specifying the name of the transposed module. The
        default name is constructed by appending "_transpose"
        to `self.module_name`.
      output_channels: Optional iterable of numbers of output channels.
      kernel_shapes: Optional iterable of kernel sizes. The default value is
        constructed by reversing `self.kernel_shapes`.
      strides: Optional iterable of kernel strides. The default value is
        constructed by reversing `self.strides`.
      paddings: Optional iterable of padding options, either `snt.SAME` or
        `snt.VALID`; The default value is constructed by reversing
        `self.paddings`.
      activation: Optional activation op. Default value is `self.activation`.
      activate_final: Optional boolean determining if the activation and batch
        normalization, if turned on, are applied to the final layer.
      normalization_ctor: Constructor to return a callable which will perform
        normalization at each layer. Defaults to None / no normalization.
        Examples of what could go here: `snt.BatchNormV2`, `snt.LayerNorm`. If
        a string is provided, importlib is used to convert the string to a
        callable, so either `snt.LayerNorm` or `"snt.LayerNorm"` can be
        provided.
      normalization_kwargs: kwargs to be provided to `normalization_ctor` when
        it is called.
      normalize_final: Whether to apply normalization after the final conv
        layer. Default is to take the value of activate_final.
      initializers: Optional dict containing ops to initialize the filters of
        the whole network (with key 'w') or biases (with key 'b'). The default
        value is `self.initializers`.
      partitioners: Optional dict containing partitioners to partition
        weights (with key 'w') or biases (with key 'b'). The default value is
        `self.partitioners`.
      regularizers: Optional dict containing regularizers for the filters of the
        whole network (with key 'w') or biases (with key 'b'). The default is
        `self.regularizers`.
      use_bias: Optional boolean or iterable of booleans determining whether to
        include bias parameters in the convolutional layers. Default
        is constructed by reversing `self.use_bias`.
      data_format: Optional string, one of "NCHW" or "NHWC". Specifies whether
        the channel dimension of the input and output is the last dimension.
        Default is `self._data_format`.
    Returns:
      Matching transposed module.

    Raises:
      ValueError: If output_channels is specified and its length does not match
        the number of layers.
    """

    if data_format is None:
      data_format = self._data_format

    if output_channels is None:
      output_channels = []
      channel_dim = -1 if data_format == DATA_FORMAT_NHWC else 1
      for layer in reversed(self._layers):
        output_channels.append(lambda l=layer: l.input_shape[channel_dim])

    elif len(output_channels) != len(self._layers):
      # Note that we only have to do this check for the output channels. Any
      # other inconsistencies will be picked up by ConvNet2D.__init__.
      raise ValueError("Iterable output_channels length must match the "
                       "number of layers ({}), but is {} instead.".format(
                           len(self._layers), len(output_channels)))

    if kernel_shapes is None:
      kernel_shapes = reversed(self.kernel_shapes)

    if strides is None:
      strides = reversed(self.strides)

    if paddings is None:
      paddings = reversed(self.paddings)

    if activation is None:
      activation = self.activation

    if activate_final is None:
      activate_final = self.activate_final

    if normalization_ctor is None:
      normalization_ctor = self.normalization_ctor

    if normalization_kwargs is None:
      normalization_kwargs = self.normalization_kwargs

    if normalize_final is None:
      normalize_final = self.normalize_final

    if initializers is None:
      initializers = self.initializers

    if partitioners is None:
      partitioners = self.partitioners

    if regularizers is None:
      regularizers = self.regularizers

    if use_bias is None:
      use_bias = reversed(self.use_bias)

    if name is None:
      name = self.module_name + "_transpose"

    return transpose_constructor(
        output_channels=output_channels,
        kernel_shapes=kernel_shapes,
        strides=strides,
        paddings=paddings,
        activation=activation,
        activate_final=activate_final,
        normalization_ctor=normalization_ctor,
        normalization_kwargs=normalization_kwargs,
        normalize_final=normalize_final,
        initializers=initializers,
        partitioners=partitioners,
        regularizers=regularizers,
        use_bias=use_bias,
        data_format=data_format,
        name=name)

Response: Returns a new version of the network with the specified parameters.


Instruction:
Implement a Python 3 function for
returning a new version of the network with the same name and parameters.

Response:
def transpose(self,
                name=None,
                output_channels=None,
                kernel_shapes=None,
                strides=None,
                paddings=None,
                activation=None,
                activate_final=None,
                normalization_ctor=None,
                normalization_kwargs=None,
                normalize_final=None,
                initializers=None,
                partitioners=None,
                regularizers=None,
                use_batch_norm=None,
                use_bias=None,
                batch_norm_config=None,
                data_format=None,
                custom_getter=None):
    """Returns transposed version of this network.

    Args:
      name: Optional string specifying the name of the transposed module. The
        default name is constructed by appending "_transpose"
        to `self.module_name`.
      output_channels: Optional iterable of numbers of output channels.
      kernel_shapes: Optional iterable of kernel sizes. The default value is
        constructed by reversing `self.kernel_shapes`.
      strides: Optional iterable of kernel strides. The default value is
        constructed by reversing `self.strides`.
      paddings: Optional iterable of padding options, either `snt.SAME` or
        `snt.VALID`; The default value is constructed by reversing
        `self.paddings`.
      activation: Optional activation op. Default value is `self.activation`.
      activate_final: Optional boolean determining if the activation and batch
        normalization, if turned on, are applied to the final layer.
      normalization_ctor: Constructor to return a callable which will perform
        normalization at each layer. Defaults to None / no normalization.
        Examples of what could go here: `snt.BatchNormV2`, `snt.LayerNorm`. If
        a string is provided, importlib is used to convert the string to a
        callable, so either `snt.LayerNorm` or `"snt.LayerNorm"` can be
        provided.
      normalization_kwargs: kwargs to be provided to `normalization_ctor` when
        it is called.
      normalize_final: Whether to apply normalization after the final conv
        layer. Default is to take the value of activate_final.
      initializers: Optional dict containing ops to initialize the filters of
        the whole network (with key 'w') or biases (with key 'b'). The default
        value is `self.initializers`.
      partitioners: Optional dict containing partitioners to partition
        weights (with key 'w') or biases (with key 'b'). The default value is
        `self.partitioners`.
      regularizers: Optional dict containing regularizers for the filters of the
        whole network (with key 'w') or biases (with key 'b'). The default is
        `self.regularizers`.
      use_batch_norm: Optional boolean determining if batch normalization is
        applied after convolution. The default value is `self.use_batch_norm`.
      use_bias: Optional boolean or iterable of booleans determining whether to
        include bias parameters in the convolutional layers. Default
        is constructed by reversing `self.use_bias`.
      batch_norm_config: Optional mapping of additional configuration for the
        `snt.BatchNorm` modules. Default is `self.batch_norm_config`.
      data_format: Optional string, one of "NCHW" or "NHWC". Specifies whether
        the channel dimension of the input and output is the last dimension.
        Default is `self._data_format`.
      custom_getter: Callable or dictionary of callables to use as
        custom getters inside the module. If a dictionary, the keys
        correspond to regexes to match variable names. See the
        `tf.get_variable` documentation for information about the
        custom_getter API.

    Returns:
      Matching `ConvNet2DTranspose` module.

    Raises:
      ValueError: If output_channels is specified and its length does not match
        the number of layers.
      ValueError: If the given data_format is not a supported format ("NHWC" or
        "NCHW").
      NotImplementedError: If the convolutions are dilated.
    """
    for rate in self._rates:
      if rate != 1:
        raise NotImplementedError("Transpose dilated convolutions "
                                  "are not supported")
    output_shapes = []
    if data_format is None:
      data_format = self._data_format
    if data_format == DATA_FORMAT_NHWC:
      start_dim, end_dim = 1, -1
    elif data_format == DATA_FORMAT_NCHW:
      start_dim, end_dim = 2, 4
    else:
      raise ValueError("Invalid data_format {:s}. Allowed formats "
                       "{}".format(data_format, SUPPORTED_2D_DATA_FORMATS))

    if custom_getter is None and self._custom_getter is not None:
      tf.logging.warning(
          "This convnet was constructed with a custom getter, but the "
          "`transpose` method was not given any. The transposed ConvNet will "
          "not be using any custom_getter.")

    for layer in reversed(self._layers):
      output_shapes.append(lambda l=layer: l.input_shape[start_dim:end_dim])
    transpose_constructor = functools.partial(ConvNet2DTranspose,
                                              output_shapes=output_shapes,
                                              custom_getter=custom_getter)

    return self._transpose(
        transpose_constructor=transpose_constructor,
        name=name,
        output_channels=output_channels,
        kernel_shapes=kernel_shapes,
        strides=strides,
        paddings=paddings,
        activation=activation,
        activate_final=activate_final,
        normalization_ctor=normalization_ctor,
        normalization_kwargs=normalization_kwargs,
        normalize_final=normalize_final,
        initializers=initializers,
        partitioners=partitioners,
        regularizers=regularizers,
        use_bias=use_bias,
        data_format=data_format)


Instruction:
Can you implement a function in Python 3 that
returns a new version of the network with the same name and parameters.

Response:
def transpose(self,
                name=None,
                output_channels=None,
                kernel_shapes=None,
                strides=None,
                paddings=None,
                activation=None,
                activate_final=None,
                normalization_ctor=None,
                normalization_kwargs=None,
                normalize_final=None,
                initializers=None,
                partitioners=None,
                regularizers=None,
                use_batch_norm=None,
                use_bias=None,
                batch_norm_config=None,
                data_format=None,
                custom_getter=None):
    """Returns transposed version of this network.

    Args:
      name: Optional string specifying the name of the transposed module. The
        default name is constructed by appending "_transpose"
        to `self.module_name`.
      output_channels: Optional iterable of numbers of output channels.
      kernel_shapes: Optional iterable of kernel sizes. The default value is
        constructed by reversing `self.kernel_shapes`.
      strides: Optional iterable of kernel strides. The default value is
        constructed by reversing `self.strides`.
      paddings: Optional iterable of padding options, either `snt.SAME` or
        `snt.VALID`; The default value is constructed by reversing
        `self.paddings`.
      activation: Optional activation op. Default value is `self.activation`.
      activate_final: Optional boolean determining if the activation and batch
        normalization, if turned on, are applied to the final layer.
      normalization_ctor: Constructor to return a callable which will perform
        normalization at each layer. Defaults to None / no normalization.
        Examples of what could go here: `snt.BatchNormV2`, `snt.LayerNorm`. If
        a string is provided, importlib is used to convert the string to a
        callable, so either `snt.LayerNorm` or `"snt.LayerNorm"` can be
        provided.
      normalization_kwargs: kwargs to be provided to `normalization_ctor` when
        it is called.
      normalize_final: Whether to apply normalization after the final conv
        layer. Default is to take the value of activate_final.
      initializers: Optional dict containing ops to initialize the filters of
        the whole network (with key 'w') or biases (with key 'b'). The default
        value is `self.initializers`.
      partitioners: Optional dict containing partitioners to partition
        weights (with key 'w') or biases (with key 'b'). The default value is
        `self.partitioners`.
      regularizers: Optional dict containing regularizers for the filters of the
        whole network (with key 'w') or biases (with key 'b'). The default is
        `self.regularizers`.
      use_batch_norm: Optional boolean determining if batch normalization is
        applied after convolution. The default value is `self.use_batch_norm`.
      use_bias: Optional boolean or iterable of booleans determining whether to
        include bias parameters in the convolutional layers. Default
        is constructed by reversing `self.use_bias`.
      batch_norm_config: Optional mapping of additional configuration for the
        `snt.BatchNorm` modules. Default is `self.batch_norm_config`.
      data_format: Optional string, one of "NCHW" or "NHWC". Specifies whether
        the channel dimension of the input and output is the last dimension.
        Default is `self._data_format`.
      custom_getter: Callable or dictionary of callables to use as
        custom getters inside the module. If a dictionary, the keys
        correspond to regexes to match variable names. See the
        `tf.get_variable` documentation for information about the
        custom_getter API.

    Returns:
      Matching `ConvNet2D` module.

    Raises:
      ValueError: If output_channels is specified and its length does not match
        the number of layers.
    """
    if use_batch_norm is not None:
      if normalization_ctor is not None or normalization_kwargs is not None:
        raise ValueError(
            "If use_batch_norm is specified, normalization_ctor and "
            "normalization_kwargs must not be.")
      if use_batch_norm:
        normalization_ctor = batch_norm.BatchNorm
      else:
        normalization_ctor = None
      normalization_kwargs = batch_norm_config

    if custom_getter is None and self._custom_getter is not None:
      tf.logging.warning(
          "This convnet was constructed with a custom getter, but the "
          "`transpose` method was not given any. The transposed ConvNet will "
          "not be using any custom_getter.")

    transpose_constructor = functools.partial(ConvNet2D,
                                              custom_getter=custom_getter)

    return self._transpose(
        transpose_constructor=transpose_constructor,
        name=name,
        output_channels=output_channels,
        kernel_shapes=kernel_shapes,
        strides=strides,
        paddings=paddings,
        activation=activation,
        activate_final=activate_final,
        normalization_ctor=normalization_ctor,
        normalization_kwargs=normalization_kwargs,
        normalize_final=normalize_final,
        initializers=initializers,
        partitioners=partitioners,
        regularizers=regularizers,
        use_bias=use_bias,
        data_format=data_format)


Instruction:
Create a Python 3 function to
load the data or reads it from cache.

Response:
def _get_raw_data(subset):
  """Loads the data or reads it from cache."""
  raw_data = _LOADED.get(subset)
  if raw_data is not None:
    return raw_data, _LOADED["vocab"]
  else:
    train_data, valid_data, test_data, vocab = ptb_reader.ptb_raw_data(
        FLAGS.data_path)
    _LOADED.update({
        "train": np.array(train_data),
        "valid": np.array(valid_data),
        "test": np.array(test_data),
        "vocab": vocab
    })
    return _LOADED[subset], vocab


Instruction:
Make a summary of the following Python 3 code
def custom_scale_mixture_prior_builder(getter, name, *args, **kwargs):
  # This specific prior formulation doesn't need any of the arguments forwarded
  # from `get_variable`.
  del getter
  del name
  del args
  del kwargs
  return CustomScaleMixture(
      FLAGS.prior_pi, FLAGS.prior_sigma1, FLAGS.prior_sigma2)

Response: A custom scale - mixture prior of Fortunato et al.


Instruction:
Given the following Python 3 function, write the documentation
def lstm_posterior_builder(getter, name, *args, **kwargs):
  del args
  parameter_shapes = tfp.distributions.Normal.param_static_shapes(
      kwargs["shape"])

  # The standard deviation of the scale mixture prior.
  prior_stddev = np.sqrt(
      FLAGS.prior_pi * np.square(FLAGS.prior_sigma1) +
      (1 - FLAGS.prior_pi) * np.square(FLAGS.prior_sigma2))

  loc_var = getter(
      "{}/posterior_loc".format(name),
      shape=parameter_shapes["loc"],
      initializer=kwargs.get("initializer"),
      dtype=tf.float32)
  scale_var = getter(
      "{}/posterior_scale".format(name),
      initializer=tf.random_uniform(
          minval=np.log(np.exp(prior_stddev / 4.0) - 1.0),
          maxval=np.log(np.exp(prior_stddev / 2.0) - 1.0),
          dtype=tf.float32,
          shape=parameter_shapes["scale"]))
  return tfp.distributions.Normal(
      loc=loc_var,
      scale=tf.nn.softplus(scale_var) + 1e-5,
      name="{}/posterior_dist".format(name))

Response: A builder for a particular diagonal gaussian posterior.