def ensure_list(x):
def linear(args, output_size, bias, bias_start=0.0, initializer=None,
           scope=None):
def _linear(args, output_size, bias, bias_start=0.0):
def _find_var_init_param(var_init_params, var_name, key, default):
def find_initializer(var_init_params, name, default_initializer):
def linear_v2(args, output_size, bias, rank=None, scope=None,   # pylint: disable=missing-docstring
              weight_name=_WEIGHTS_VARIABLE_NAME,
              bias_name=_BIAS_VARIABLE_NAME,
              var_init_params=None):
def layer_norm(x, reduction_indices, epsilon=1e-9, gain=None, bias=None,
               per_element=True, scope=None):
def get_sparse_variable(name, indices, shape, dtype=None, trainable=True,
                        initializer=None, partitioner=None, regularizer=None):
def _random_index(shape):
def _all_indices(shape):  # pylint: disable=missing-docstring
  def add_indices(shape, depth, index):
def sparse_random_indices(ratio, shape):
def get_optimizer(optimizer_type):
  def build(learning_rate, config):
def _extract_dict_from_config(config, prefix, keys):
def repeat(iterable, n):
def is_var_in_scope(var, scopes):
def trainable_vars_in_scope(name):
def find_var(name, vars_=None):
def group_vars_by_scope(scopes, vars_=None, log=False):
def _order_grouped_vars(var_groups):
def gradients_for_var_group(var_groups, gradients, name):
def trainable_initial_state(batch_size, state_size, initial_state_init=None):
  def create_one(i, size):
def map_nested(fn, x):
def create_grads(optimizer, loss, scopes, num_expected_missing_gradients=0):
def clip_gradients_in_scope(grads_and_vars, scope, max_grad_norm):
def cv_splits(seq, k=10):
def escape_summary_name(name):
def summaries_for_trainables():
def log_scalar_summaries(summary_proto):
def count_trainables(scopes=('',)):
def log_trainables(scopes=('',)):
def round_to_int(x):
def orthogonal_initializer(gain=1.0, dtype=tf.float32):
  def initialize(shape, dtype=dtype, partition_info=None, gain=gain):
def random_mask(shape, k):
def random_mask2(shape, k):
def make_low_rank_factorization_initializer(shape, rank):
def low_rank_factorization(name, shape, rank, initializer=None,
                           trainable=True, collections=None):
  def __init__(self, name='serialized', initializer='{}'):
  def store(self, obj):
  def retrieve(self):
  def variables(self):
def mixture_of_softmaxes(x, k, e, to_logits):
def expand_tile(tensor, n, name=None):
def get_batched_variable(name, runtime_batch_size, shape=None,
                         dtype=tf.float32, initializer=None,
                         trainable=True):
def create_mask(lengths, max_length):
def mask_from_lengths(lengths, max_length=None, dtype=None, name=None):
def compute_lengths(symbols_list, eos_symbol, name=None,
                    dtype=tf.int64):
def seq_softmax_cross_entropy_with_logits(logits, labels, lengths,
                                          max_length=None, reduce_sum=True,
                                          name=None):
def variance_scaling_initializer(scale=2.0, mode='fan_in',
                                 distribution='truncated_normal',
                                 mean=0.0, seed=None, dtype=tf.float32):
  def _initializer(shape, dtype=dtype, partition_info=None):
  def __init__(self, fn, output_size, reuse=None):
  def state_size(self):
  def output_size(self):
  def zero_state(self, batch_size, dtype):
  def call(self, inputs, state):
_BIAS_VARIABLE_NAME = 'bias'
_WEIGHTS_VARIABLE_NAME = 'kernel'
class TFSerializer(object):
class FNCell(tf.nn.rnn_cell.RNNCell): class AbstractTiledLinear(snt_base.AbstractModule):
class TiledLinear(AbstractTiledLinear):
class LayerNormedTiledLinear(AbstractTiledLinear):
class SparseTiledLinear(AbstractTiledLinear):
class OverlaidTiledLinear(TiledLinear):
  def __init__(self, input_name_and_sizes, output_name_and_sizes,
               var_init_params=None, name='tiled_linear'):
  def _canonicalize_input_name_and_sizes(self, name_and_sizes):
  def _check_output_name_and_sizes(self, name_and_sizes):
  def _check_var_init_params(self, var_init_params):
  def _check_dtype(self, inputs, previous_dtype):
  def valid_var_init_param_keys(self):
  def _find_var_init_param(self, var_name, key, default):
  def _get_variable(self, name, shape,
                    initializer=None,
                    default_initializer=None, default_partitioner=None,
                    default_regularizer=None):
  def _declared_input_sizes(self):
  def _inferred_input_sizes(self, inputs):
  def _merge_input_sizes(self, inputs):
  def _merged_input_name_and_sizes(self):
  def _output_name_and_sizes(self):
  def _build(self, inputs):
  def __init__(self, input_name_and_sizes, output_name_and_sizes,
               var_init_params=None, name='tiled_linear'):
  def _ensure_weights(self):
      def tiled_initializer(shape, dtype=self._dtype, partition_info=None):
  def _ensure_biases(self):
  def _build_tiled_linear(self, inputs, input_name_and_sizes,
                          output_name_and_sizes, add_bias):
    def split_output(output):
  def _build_tiled_linear(self, inputs, input_name_and_sizes,
                          output_name_and_sizes, add_bias):
    def make_weights_for_inputs():
    def make_biases():
    def split_output(output):
  def __init__(self, input_name_and_sizes, output_name_and_sizes,
               var_init_params=None, name='sparse_tiled_linear'):
  def _find_or_create_sparse_indices(self, name, shape):
  def _find_or_create_sparse_variable(self, name, sparse_indices, shape,
                                      initializer=None, partitioner=None,
                                      regularizer=None):
  def valid_var_init_param_keys(self):
  def _get_variable(self, name, shape,
                    default_initializer=None, default_partitioner=None,
                    default_regularizer=None, sparse_indices=None):
  def _build_tiled_linear(self, inputs, input_name_and_sizes,
                          output_name_and_sizes, add_bias):
  def __init__(self, *args, **kwargs):
  def _get_variable(self, name, shape,
                    default_initializer=None, default_partitioner=None,
                    default_regularizer=None):
  def _low_rank_matrix(self, name, rank=None, shape=None,
                       initializer=None, trainable=True):
  def valid_var_init_param_keys(self): class NASCell(tf.nn.rnn_cell.RNNCell):
  """Neural Architecture Search (NAS) recurrent network cell.

  This implements the recurrent cell from the paper:

    https://arxiv.org/abs/1611.01578

  Barret Zoph and Quoc V. Le.
  "Neural Architecture Search with Reinforcement Learning" Proc. ICLR 2017.

  The class uses an optional projection layer.
  """

  def __init__(self, num_units, num_proj=None,
               use_biases=False, reuse=None,
               initializer=None,
               input_transform=None,
               state_transform=None,
               update_transform=None):
    """Initialize the parameters for a NAS cell.

    Args:
      num_units: int, The number of units in the NAS cell
      num_proj: (optional) int, The output dimensionality for the projection
        matrices.  If None, no projection is performed.
      use_biases: (optional) bool, If True then use biases within the cell. This
        is False by default.
      reuse: (optional) Python boolean describing whether to reuse variables
        in an existing scope.  If not `True`, and the existing scope already has
        the given variables, an error is raised.
      initializer: Initializer for the variables.
      input_transform: None, or a function of one argument that
        massages the input in some way. For example, variational
        dropout can be implemted by passing a Dropout object here.
      state_transform: Similar to input_transform, this is
        applied to the recurrent state.
      update_transform: Similar to input_transform, this is
        applied to the proposed update ('j').
    """
    super(NASCell, self).__init__(_reuse=reuse)
    self._num_units = num_units
    self._num_proj = num_proj
    self._use_biases = use_biases
    self._reuse = reuse

    if num_proj is not None:
      self._state_size = tf.nn.rnn_cell.LSTMStateTuple(num_units, num_proj)
      self._output_size = num_proj
    else:
      self._state_size = tf.nn.rnn_cell.LSTMStateTuple(num_units, num_units)
      self._output_size = num_units
    self._initializer = initializer
    self._input_transform = input_transform
    self._state_transform = state_transform
    assert update_transform is None

  @property
  def state_size(self):
    return self._state_size

  @property
  def output_size(self):
    return self._output_size

  def call(self, inputs, state):
    """Run one step of NAS Cell.

    Args:
      inputs: input Tensor, 2D, batch x num_units.
      state: This must be a tuple of state Tensors, both `2-D`, with column
        sizes `c_state` and `m_state`.

    Returns:
      A tuple containing:
      - A `2-D, [batch x output_dim]`, Tensor representing the output of the
        NAS Cell after reading `inputs` when previous state was `state`.
        Here output_dim is:
           num_proj if num_proj was set,
           num_units otherwise.
      - Tensor(s) representing the new state of NAS Cell after reading `inputs`
        when the previous state was `state`.  Same type and shape(s) as `state`.

    Raises:
      ValueError: If input size cannot be inferred from inputs via
        static shape inference.
    """
    sigmoid = tf.sigmoid
    tanh = tf.tanh
    relu = tf.nn.relu

    num_proj = self._num_units if self._num_proj is None else self._num_proj

    def maybe_transform(transform, x):
      if transform is None:
        return x
      else:
        return transform(x)

    (c_prev, m_prev) = state
    m_prev = maybe_transform(self._state_transform, m_prev)

    dtype = inputs.dtype
    input_size = inputs.get_shape().with_rank(2)[1]
    inputs = maybe_transform(self._input_transform, inputs)
    if input_size.value is None:
      raise ValueError("Could not infer input size from inputs.get_shape()[-1]")
    # Variables for the NAS cell. W_m is all matrices multiplying the
    # hiddenstate and W_inputs is all matrices multiplying the inputs.
    concat_w_m = tf.get_variable(
        "recurrent_kernel", [num_proj, 8 * self._num_units],
        initializer=self._initializer, dtype=dtype)
    concat_w_inputs = tf.get_variable(
        "kernel", [input_size.value, 8 * self._num_units],
        initializer=self._initializer, dtype=dtype)

    m_matrix = tf.matmul(m_prev, concat_w_m)
    inputs_matrix = tf.matmul(inputs, concat_w_inputs)

    if self._use_biases:
      b = tf.get_variable(
          "bias",
          shape=[8 * self._num_units],
          initializer=tf.zeros_initializer(),
          dtype=dtype)
      m_matrix = tf.nn.bias_add(m_matrix, b)

    # The NAS cell branches into 8 different splits for both the hiddenstate
    # and the input
    m_matrix_splits = tf.split(axis=1, num_or_size_splits=8,
                               value=m_matrix)
    inputs_matrix_splits = tf.split(axis=1, num_or_size_splits=8,
                                    value=inputs_matrix)

    # First layer
    layer1_0 = sigmoid(inputs_matrix_splits[0] + m_matrix_splits[0])
    layer1_1 = relu(inputs_matrix_splits[1] + m_matrix_splits[1])
    layer1_2 = sigmoid(inputs_matrix_splits[2] + m_matrix_splits[2])
    layer1_3 = relu(inputs_matrix_splits[3] * m_matrix_splits[3])
    layer1_4 = tanh(inputs_matrix_splits[4] + m_matrix_splits[4])
    layer1_5 = sigmoid(inputs_matrix_splits[5] + m_matrix_splits[5])
    layer1_6 = tanh(inputs_matrix_splits[6] + m_matrix_splits[6])
    layer1_7 = sigmoid(inputs_matrix_splits[7] + m_matrix_splits[7])

    # Second layer
    l2_0 = tanh(layer1_0 * layer1_1)
    l2_1 = tanh(layer1_2 + layer1_3)
    l2_2 = tanh(layer1_4 * layer1_5)
    l2_3 = sigmoid(layer1_6 + layer1_7)

    # Inject the cell
    l2_0 = tanh(l2_0 + c_prev)

    # Third layer
    l3_0_pre = l2_0 * l2_1
    new_c = l3_0_pre  # create new cell
    l3_0 = l3_0_pre
    l3_1 = tanh(l2_2 + l2_3)

    # Final layer
    new_m = tanh(l3_0 * l3_1)

    # Projection layer if specified
    if self._num_proj is not None:
      concat_w_proj = tf.get_variable(
          "projection_weights", [self._num_units, self._num_proj],
          dtype)
      new_m = tf.matmul(new_m, concat_w_proj)

    new_state = tf.nn.rnn_cell.LSTMStateTuple(new_c, new_m)
    return new_m, new_state class TiledLSTMCell(tf.nn.rnn_cell.RNNCell):
  """An LSTM cell with tiled connections.

  Supports various connectivity patterns such as the vanilla, dense
  TiledLinear, and also SparseTiledLinear, LayerNormedTiledLinear.
  """

  def __init__(
      self, num_units,
      use_peepholes=False, cell_clip=None,
      initializer=None, num_proj=None,
      feature_mask_rounds=0,
      feature_mask_rank=0,
      tie_gates=False,
      cap_input_gate=True,
      layer_norm=False,
      activation=tf.tanh,
      input_transform=None,
      state_transform=None,
      update_transform=None,
      tiled_linear_class=None,
      tiled_linear_var_init_params=None):
    """Initialize the parameters of a single LSTM layer.

    Args:
      num_units: int, The number of hidden units in the layer.
      use_peepholes: bool, set True to enable diagonal/peephole connections
        (non implemented).
      cell_clip: (optional) A float value, if provided the cell state
        is clipped to be in the [-cell_clip, cell_clip] range prior to
        the cell output activation.
      initializer: (optional) The default initializer to use for the
        weight and projection matrices.
      num_proj: (optional) int, The output size of the non-linear
        transformation (usually `h`) of the cell state (`c`). If None,
        no projection is performed and `h=tanh(c)`. If provided, then
        `h` is `tanh(c)` projected to `num_proj` dimensions.
      feature_mask_rounds: Gate the input and the state before they are used for
        calculating all the other stuff (i.e. i, j, o, f). This allows input
        features to be reweighted based on the state, and state features to be
        reweighted based on the input. When feature_mask_rounds is 0, there is
        no extra gating in the LSTM. When 1<=, the input is gated: x *=
        2*sigmoid(affine(h))). When 2<=, the state is gated: h *=
        2*sigmoid(affine(x))). For higher number of rounds, the alternating
        gating continues.
      feature_mask_rank: If 0, the linear transforms are full rank, dense
        matrices. If >0, then the matrix representing the linear transform is
        factorized as the product of two low rank matrices ([*, rank] and [rank,
        *]). This reduces the number of parameters greatly.
      tie_gates: Whether to make the input gate one minus the forget gate.
      cap_input_gate: Whether to cap the input gate at one minus the
        forget gate (if they are not tied, of course). This ensures
        'c' is in [-1,1] and makes training easier especially in the
        early stages.
      layer_norm: Whether to use Layer Normalization.
      activation: Activation function of the inner states.
      input_transform: None, or a function of one argument that
        massages the input in some way. For example, variational
        dropout can be implemted by passing a Dropout object here.
      state_transform: Similar to input_transform, this is
        applied to the recurrent state.
      update_transform: Similar to input_transform, this is
        applied to the proposed update ('j').
      tiled_linear_class: A class such as tiled_linear.TiledLinear
        that's instantiated an unspecified number of times with the
        same tiled_linear_var_init_params but with possibly different
        inputs and outputs. If layer_norm is false, the default is
        tiled_linear.TiledLinear else it's
        tiled_linear.LayerNormedTiledLinear.
      tiled_linear_var_init_params: Passed right on to
        `tiled_linear_class` as the `var_init_params` argument.
    """
    assert not use_peepholes, 'Peepholes are not implemented in LSTMCell.'
    self._num_units = num_units
    self._use_peepholes = use_peepholes
    self._cell_clip = cell_clip
    self._initializer = initializer
    self._num_proj = num_proj
    self._feature_mask_rounds = feature_mask_rounds
    self._feature_mask_rank = feature_mask_rank
    self._tie_gates = tie_gates
    self._cap_input_gate = cap_input_gate
    self._layer_norm = layer_norm
    self._activation = activation
    self._input_transform = input_transform
    self._state_transform = state_transform
    self._update_transform = update_transform
    if tiled_linear_class is None:
      if layer_norm:
        tiled_linear_class = tiled_linear.LayerNormedTiledLinear
      else:
        tiled_linear_class = tiled_linear.TiledLinear
    self._tiled_linear_class = tiled_linear_class
    self._tiled_linear_var_init_params = tiled_linear_var_init_params
    self._tiled_linear_mod = None

    if num_proj:
      self._output_size = num_proj
    else:
      self._output_size = num_units

    self._state_size = tf.nn.rnn_cell.LSTMStateTuple(
        num_units, self._output_size)

  @property
  def state_size(self):
    return self._state_size

  @property
  def output_size(self):
    return self._output_size

  @staticmethod
  def _do_feature_masking(x, y, num_x, num_y, rounds, rank):
    for round_ in six.moves.range(rounds):
      # Even rounds correspond to input transforms. Odd rounds to state
      # transforms. Implemented this way because feature_mask_rounds=1 with a
      # single round of transforming the state does not seem to improve things
      # much. Concurrent updates were also tested, but were not an improvement
      # either.
      transforming_x = (round_ % 2 == 0)
      fm_name = 'fm_' + str(round_)
      if rank == 0:  # full rank case
        if transforming_x:
          x *= 2*tf.sigmoid(utils.linear(y, num_x, bias=True, scope=fm_name))
        else:
          y *= 2*tf.sigmoid(utils.linear(x, num_y, bias=True, scope=fm_name))
      else:  # low-rank factorization case
        if transforming_x:
          shape = [num_y, num_x]
        else:
          shape = [num_x, num_y]
        a, b = utils.low_rank_factorization(fm_name + '_weight', shape, rank)
        bias = tf.get_variable(fm_name + '_bias', shape[1],
                               initializer=tf.zeros_initializer())
        if transforming_x:
          x *= 2*tf.sigmoid(tf.matmul(tf.matmul(y, a), b) + bias)
        else:
          y *= 2*tf.sigmoid(tf.matmul(tf.matmul(x, a), b) + bias)
    return x, y

  def __call__(self, input_, state, scope=None):
    """Run one step of LSTM.

    All tensor arguments are shaped [batch_size, *].

    Args:
      input_: A tensor.
      state: An LSTMStateTuple.
      scope: VariableScope for the created subgraph; defaults to
        `LSTMCell`.

    Returns:
      A tuple containing:
      - A `2-D, [batch, output_dim]`, Tensor representing the output of the
        LSTM after one time step.
        Here output_dim is:
           - num_proj if num_proj was set,
           - num_units otherwise.
      - An LSTMStateTuple of Tensors representing the new state
        of the LSTM after one time step.

    Raises:
      ValueError: If input size cannot be inferred from `input_`
      via static shape inference.
    """
    num_units = self._num_units
    num_proj = num_units if self._num_proj is None else self._num_proj
    num_inputs = input_.get_shape().with_rank(2)[1]

    def maybe_transform(transform, x):
      if transform is None:
        return x
      else:
        return transform(x)

    with tf.variable_scope(scope or type(self).__name__,
                           initializer=self._initializer):
      # Apply transformations to the input and the recurrent state.
      transformed_input = maybe_transform(self._input_transform, input_)
      transformed_state = maybe_transform(self._state_transform, state.h)

      # Let's transform the input and state further with 'feature masking'.
      transformed_input, transformed_state = self._do_feature_masking(
          transformed_input, transformed_state,
          num_inputs, num_units,
          self._feature_mask_rounds, self._feature_mask_rank)

      inputs = [transformed_input, transformed_state]

      input_name_and_sizes = [('x', num_inputs),
                              ('h', num_proj)]
      output_name_and_sizes = [('j', num_units),
                               ('o', num_units),
                               ('f', num_units)]
      if not self._tie_gates:
        output_name_and_sizes.append(('i', num_units))

      if self._tiled_linear_mod is None:
        self._tiled_linear_mod = self._tiled_linear_class(
            input_name_and_sizes, output_name_and_sizes,
            self._tiled_linear_var_init_params)
      if self._tie_gates:
        j_pre, o_pre, f_pre = self._tiled_linear_mod(inputs)
      else:
        j_pre, o_pre, f_pre, i_pre = self._tiled_linear_mod(inputs)
      # Compute the cell state c.
      f = tf.sigmoid(f_pre)
      j = self._activation(j_pre)
      j = maybe_transform(self._update_transform, j)
      o = tf.sigmoid(o_pre)
      if self._tie_gates:
        c = f * state.c + (1-f) * j
      else:
        i = tf.sigmoid(i_pre)
        if self._cap_input_gate:
          c = f * state.c + tf.minimum(1-f, i) * j
        else:
          c = f * state.c + i * j

      if self._layer_norm:
        c2 = utils.layer_norm(c, [1], scope='ln_c')
      else:
        c2 = c

      if self._cell_clip is not None:
        # pylint: disable=invalid-unary-operand-type
        c = tf.clip_by_value(c, -self._cell_clip, self._cell_clip)
        # pylint: enable=invalid-unary-operand-type

      h = o * self._activation(c2)
      if self._num_proj is not None:
        h = utils.linear(h, self._num_proj, bias=False, scope='projection')

    return h, tf.nn.rnn_cell.LSTMStateTuple(c, h) class TiledRHNCell(tf.nn.rnn_cell.RNNCell):
  """An RHN cell with tiled connections.

  A RHN cell is like a simplified LSTM with multiple layers. See
  'Recurrent Highway Networks' by Zilly et al.
  (https://arxiv.org/abs/1607.03474). This implementation is based on
  v3 of that paper.

  Supports various connectivity patterns such as the vanilla, dense
  TiledLinear, and also SparseTiledLinear, LayerNormedTiledLinear.
  """

  def __init__(
      self, num_units, depth=1,
      use_peepholes=False, cell_clip=None,
      initializer=None,
      tie_gates=False,
      activation=tf.tanh,
      input_transform=None,
      state_transform=None,
      update_transform=None,
      tiled_linear_class=None,
      tiled_linear_var_init_params=None):
    """Initialize the parameters of a single RHN layer.

    Args:
      num_units: int, The number of hidden units in the layer.
      depth: int, The number of layers.
      use_peepholes: bool, set True to enable diagonal/peephole connections
        (non implemented).
      cell_clip: (optional) A float value, if provided the cell state
        is clipped to be in the [-cell_clip, cell_clip] range prior to
        the cell output activation.
      initializer: (optional) The initializer to use for the weight and
        projection matrices.
      tie_gates: Whether to make the input gate one minus the forget gate.
      activation: Activation function of the inner states.
      input_transform: None, or a function of one argument that
        massages the input in some way. For example, variational
        dropout can be implemted by passing a Dropout object here.
      state_transform: Similar to input_transform, this is
        applied to the recurrent state.
      update_transform: Similar to input_transform, this is
        applied to the proposed update ('h').
      tiled_linear_class: A class such as tiled_linear.TiledLinear
        that's instantiated an unspecified number of times with the
        same tiled_linear_var_init_params but with possibly different
        inputs and outputs. Defaults to tiled_linear.TiledLinear.
      tiled_linear_var_init_params: Passed right on to
        `tiled_linear_class` as the `var_init_params` argument.
    """
    assert not use_peepholes, 'Peepholes are not implemented in RHNCell.'
    self._num_units = num_units
    self._depth = depth
    self._use_peepholes = use_peepholes
    self._cell_clip = cell_clip
    self._initializer = initializer
    self._tie_gates = tie_gates
    self._activation = activation
    self._input_transform = input_transform
    self._state_transform = state_transform
    self._update_transform = update_transform
    if tiled_linear_class is None:
      tiled_linear_class = tiled_linear.TiledLinear
    self._tiled_linear_class = tiled_linear_class
    self._tiled_linear_var_init_params = tiled_linear_var_init_params
    self._tiled_linear_mods = [None]*depth

    self._output_size = num_units
    self._state_size = TiledRHNStateTuple(num_units)

  @property
  def state_size(self):
    return self._state_size

  @property
  def output_size(self):
    return self._output_size

  def __call__(self, input_, state, scope=None):
    """Run one step of RHN.

    All tensor arguments are shaped [batch_size, *].

    Args:
      input_: A tensor.
      state: An TiledRHNStateTuple.
      scope: VariableScope for the created subgraph; defaults to
        `TiledRHNCell`.

    Returns:
      A tuple containing:
      - A `2-D, [batch, num_units]`, Tensor representing the output of
        the RHN after one time step (which consists of `depth` number
        of computational steps).
      - An TiledRHNStateTuple tuple of Tensors representing the new state
        of the RHN after one time step.

    Raises:
      ValueError: If input size cannot be inferred from `input_`
      via static shape inference.
    """
    num_units = self._num_units

    def maybe_transform(transform, x):
      if transform is None:
        return x
      else:
        return transform(x)

    # Apply transformations to the input and the recurrent state.
    transformed_input = maybe_transform(self._input_transform, input_)

    # Let's figure out what the outputs are.
    output_name_and_sizes = [
        # This is the proposed update (usually 'j' in an LSTM).
        ('h', num_units),
        # Called 'carry' gate in the paper. This pretty much plays the
        # part of the forget gate of an LSTM.
        ('c', num_units)]
    if not self._tie_gates:
      # Called 'transform' gate, this is like the input gate of an
      # LSTM.
      output_name_and_sizes.append(('t', num_units))

    with tf.variable_scope(scope or type(self).__name__,
                           initializer=self._initializer):
      s = state.s
      for level in six.moves.range(self._depth):
        with tf.variable_scope('layer{}'.format(level)):
          transformed_s = maybe_transform(self._state_transform, s)
          if level == 0:
            inputs = [transformed_input, transformed_s]
            input_name_and_sizes = [
                ('x', transformed_input.get_shape().with_rank(2)[1]),
                # This is the raw cell state. Unlike in an LSTM this
                # is not passed through any non-linearity.
                ('s', num_units)]
          else:
            inputs = [transformed_s]
            input_name_and_sizes = [('s', num_units)]
          if self._tiled_linear_mods[level] is None:
            self._tiled_linear_mods[level] = self._tiled_linear_class(
                input_name_and_sizes, output_name_and_sizes,
                self._tiled_linear_var_init_params)
          if self._tie_gates:
            h_pre, c_pre = self._tiled_linear_mods[level](inputs)
          else:
            h_pre, c_pre, t_pre = self._tiled_linear_mods[level](inputs)
          # Compute the cell state s.
          c = tf.sigmoid(c_pre)
          h = self._activation(h_pre)
          h = maybe_transform(self._update_transform, h)
          if self._tie_gates:
            t = 1 - c
          else:
            t = tf.sigmoid(t_pre)
          s = c * s + t * h

          if self._cell_clip is not None:
            # pylint: disable=invalid-unary-operand-type
            s = tf.clip_by_value(s, -self._cell_clip, self._cell_clip)
            # pylint: enable=invalid-unary-operand-type

    return s, TiledRHNStateTuple(s) class ResMultiRNNCell(tf.nn.rnn_cell.RNNCell):
  """RNN cell composed sequentially of multiple simple cells."""

  def __init__(self, cells, state_is_tuple=True):
    """Create a RNN cell composed sequentially of a number of RNNCells.

    Args:
      cells: list of RNNCells that will be composed in this order.
      state_is_tuple: If True, accepted and returned states are n-tuples, where
        `n = len(cells)`.  If False, the states are all
        concatenated along the column axis.  This latter behavior will soon be
        deprecated.

    Raises:
      ValueError: if cells is empty (not allowed), or at least one of the cells
        returns a state tuple but the flag `state_is_tuple` is `False`.
    """
    if not cells:
      raise ValueError("Must specify at least one cell for ResMultiRNNCell.")
    if not nest.is_sequence(cells):
      raise TypeError(
          "cells must be a list or tuple, but saw: %s." % cells)

    self._cells = cells
    self._state_is_tuple = state_is_tuple
    if not state_is_tuple:
      if any(nest.is_sequence(c.state_size) for c in self._cells):
        raise ValueError("Some cells return tuples of states, but the flag "
                         "state_is_tuple is not set.  State sizes are: %s"
                         % str([c.state_size for c in self._cells]))

  @property
  def state_size(self):
    if self._state_is_tuple:
      return tuple(cell.state_size for cell in self._cells)
    else:
      return sum([cell.state_size for cell in self._cells])

  @property
  def output_size(self):
    return self._cells[-1].output_size

  def __call__(self, inputs, state, scope=None):
    """Run this multi-layer cell on inputs, starting from state."""
    with tf.variable_scope(scope or "res_multi_rnn_cell"):
      cur_state_pos = 0
      cur_inp = inputs
      new_states = []
      for i, cell in enumerate(self._cells):
        with tf.variable_scope("cell_%d" % i):
          if self._state_is_tuple:
            if not nest.is_sequence(state):
              raise ValueError(
                  "Expected state to be a tuple of length %d, but received: %s"
                  % (len(self.state_size), state))
            cur_state = state[i]
          else:
            cur_state = tf.slice(
                state, [0, cur_state_pos], [-1, cell.state_size])
            cur_state_pos += cell.state_size
          cur_inp2, new_state = cell(cur_inp, cur_state)
          if i == 0:
            cur_inp = cur_inp2
          else:
            cur_inp = cur_inp + cur_inp2
          new_states.append(new_state)
    new_states = (tuple(new_states) if self._state_is_tuple else
                  tf.concat(new_states, 1))
    return cur_inp, new_states class SkipMultiRNNCell(tf.nn.rnn_cell.RNNCell):
  """RNN cell composed sequentially of multiple simple cells."""

  def __init__(self, cells, state_is_tuple=True):
    """Create a RNN cell composed sequentially of a number of RNNCells.

    Args:
      cells: list of RNNCells that will be composed in this order.
      state_is_tuple: If True, accepted and returned states are n-tuples, where
        `n = len(cells)`.  If False, the states are all
        concatenated along the column axis.  This latter behavior will soon be
        deprecated.

    Raises:
      ValueError: if cells is empty (not allowed), or at least one of the cells
        returns a state tuple but the flag `state_is_tuple` is `False`.
    """
    if not cells:
      raise ValueError("Must specify at least one cell for SkipMultiRNNCell.")
    if not nest.is_sequence(cells):
      raise TypeError(
          "cells must be a list or tuple, but saw: %s." % cells)

    self._cells = cells
    self._state_is_tuple = state_is_tuple
    if not state_is_tuple:
      if any(nest.is_sequence(c.state_size) for c in self._cells):
        raise ValueError("Some cells return tuples of states, but the flag "
                         "state_is_tuple is not set.  State sizes are: %s"
                         % str([c.state_size for c in self._cells]))

  @property
  def state_size(self):
    if self._state_is_tuple:
      return tuple(cell.state_size for cell in self._cells)
    else:
      return sum([cell.state_size for cell in self._cells])

  @property
  def output_size(self):
    return self._cells[-1].output_size

  def __call__(self, inputs, state, scope=None):
    """Run this multi-layer cell on inputs, starting from state."""
    output = None
    with tf.variable_scope(scope or "skip_multi_rnn_cell"):
      cur_state_pos = 0
      cur_inp = inputs
      new_states = []
      for i, cell in enumerate(self._cells):
        with tf.variable_scope("cell_%d" % i):
          if self._state_is_tuple:
            if not nest.is_sequence(state):
              raise ValueError(
                  "Expected state to be a tuple of length %d, but received: %s"
                  % (len(self.state_size), state))
            cur_state = state[i]
          else:
            cur_state = tf.slice(
                state, [0, cur_state_pos], [-1, cell.state_size])
            cur_state_pos += cell.state_size
          cur_inp, new_state = cell(cur_inp, cur_state)
          new_states.append(new_state)
          if output is None:
            output = cur_inp
          else:
            output += cur_inp
    new_states = (tuple(new_states) if self._state_is_tuple else
                  tf.concat(new_states, 1))
    return output, new_states class DirichletDropout(snt_base.AbstractModule):
  """Possibly variational dropout."""

  def __init__(self, keep_prob, share_mask=True, scaler=1.0, name='dropout'):
    super(DirichletDropout, self).__init__(name=name)
    self._keep_prob = keep_prob
    self._keep_mask = None
    self._share_mask = share_mask
    self._scaler = scaler

  def _ensure_keep_mask(self, x):
    if self._keep_mask is None or not self._share_mask:
      shape = tf.shape(x)
      k = shape[1]
      # To make this class a drop-in replacement for bernoulli dropout we
      # paramaterize it with keep_prob. Set alpha of the dirichlet so that the
      # variance is equal to the variance of the bernoulli with p=keep_prob
      # divided by keep_prob.
      # Now the variance of the dirichlet with k equal alphas is
      # (k-1)/(k^2(k*alpha+1). Solve that for alpha.
      kf = tf.cast(k, tf.float32)
      alpha = self._keep_prob * (kf - 1.0) / ((1-self._keep_prob)*kf) - 1.0/kf
      dist = tfp.distributions.Dirichlet(tf.ones(shape=k) * alpha)
      assert (dist.reparameterization_type ==
              tfp.distributions.FULLY_REPARAMETERIZED)
      # The E[dir(alpha)] = 1/k for all elements, but we want the expectation to
      # be keep_prob, hence the multiplication.
      self._keep_mask = kf * dist.sample(shape[0])
      self._keep_mask.set_shape(x.get_shape())
    return self._keep_mask

  def _build(self, x):
    if contrib_util.constant_value(self._keep_prob) == 1:
      return x
    else:
      return tf.cond(tf.equal(self._keep_prob, 1.0),
                     lambda: x,
                     lambda: x * self._ensure_keep_mask(x)) class DriftingDropout(snt_base.AbstractModule):
  """Dropout with gradually changing mask."""

  def __init__(self, keep_prob, flip_prob=0.0, scaler=1.0, name='dropout'):
    super(DriftingDropout, self).__init__(name=name)
    self._keep_prob = keep_prob
    self._flip_prob = flip_prob
    self._scaler = scaler
    self._time_step = 0

  def _build(self, x, state):
    prev_keep_mask = state
    shape = tf.shape(x)
    noise = tf.random_uniform(shape, dtype=x.dtype)
    other_mask = tf.floor(self._keep_prob + noise)
    choice_noise = tf.random_uniform(shape, dtype=x.dtype)
    choice = tf.less(choice_noise, self._flip_prob)
    # KLUDGE(melisgl): The client has to pass the last keep_mask from
    # a batch to the next so the mask may end up next to some
    # recurrent cell state. This state is often zero at the beginning
    # and may be periodically zeroed (per example) during training.
    # While zeroing LSTM state is okay, zeroing the dropout mask is
    # not. So instead of forcing every client to deal with this common
    # (?) case, if an all zero mask is detected, then regenerate a
    # fresh mask. This is of course a major hack and won't help with
    # learnt initial states, for example.
    sum_ = tf.reduce_sum(prev_keep_mask, 1, keepdims=True)
    is_initializing = tf.equal(sum_, 0.0)

    self._keep_mask = tf.where(tf.logical_or(choice, is_initializing),
                               other_mask,
                               prev_keep_mask)
    self._time_step += 1
    return x * self._keep_mask / self._keep_prob * self._scaler class Dropout(snt_base.AbstractModule):
  """Possibly variational dropout."""

  def __init__(self, keep_prob, share_mask=True, scaler=1.0, name='dropout'):
    super(Dropout, self).__init__(name=name)
    self._keep_prob = keep_prob
    self._keep_mask = None
    self._share_mask = share_mask
    self._scaler = scaler

  def _ensure_keep_mask(self, x):
    if self._keep_mask is None or not self._share_mask:
      shape = tf.shape(x)
      noise = tf.random_uniform(shape, dtype=x.dtype)
      self._keep_mask = (tf.floor(self._keep_prob + noise)
                         * (self._scaler / self._keep_prob))
      self._keep_mask.set_shape(x.get_shape())
    return self._keep_mask

  def _build(self, x):
    if contrib_util.constant_value(self._keep_prob) == 1:
      return x
    else:
      return x * self._ensure_keep_mask(x) class GaussianDropout(snt_base.AbstractModule):
  """Possibly variational dropout."""

  def __init__(self, keep_prob, share_mask=True, scaler=1.0, name='dropout'):
    super(GaussianDropout, self).__init__(name=name)
    self._keep_prob = keep_prob
    self._keep_mask = None
    self._share_mask = share_mask
    self._scaler = scaler

  def _ensure_keep_mask(self, x):
    if self._keep_mask is None or not self._share_mask:
      shape = tf.shape(x)
      # Calculate the stddev for the normal distribution that
      # matches the stddev of the bernoulli with p=keep_prob.
      stddev = tf.sqrt((1 - self._keep_prob) / self._keep_prob)
      self._keep_mask = tf.random_normal(shape, mean=1.0, stddev=stddev,
                                         dtype=x.dtype)
      self._keep_mask.set_shape(x.get_shape())
    return self._keep_mask

  def _build(self, x):
    if contrib_util.constant_value(self._keep_prob) == 1:
      return x
    else:
      return x * self._ensure_keep_mask(x)
import collections
import math
import six
import tensorflow.compat.v1 as tf
from absl import logging
from lamb import utils
from lamb import tiled_linear
from lamb.nascell import NASCell
from lamb.tiled_lstm import TiledLSTMCell
from lamb.tiled_rhn import TiledRHNCell
from lamb.res_multi_rnn_cell import ResMultiRNNCell
from lamb.skip_multi_rnn_cell import SkipMultiRNNCell
from lamb.dropout import DirichletDropout
from lamb.dropout import DriftingDropout
from lamb.dropout import Dropout
from lamb.dropout import GaussianDropout
from tensorflow.contrib import framework as contrib_framework
      assert overlay_rank == -1
      tiled_linear_class = tiled_linear.SparseTiledLinear
      # This is equivalent to using cell_initializer scaled by
      # 1/sparsity_ratio.
      sparse_initializer = tf.truncated_normal_initializer(
          stddev=math.sqrt(cell_init_factor /
                           sparsity_ratio /
                           # TODO(melisgl): This is off if the input
                           # embedding size is different from the hidden
                           # size.
                           hidden_size))
      init_params.update(collections.OrderedDict([
          ('W_x_.*', {'sparse_indices_sharing_key': 'W_x'}),
          ('W_h_.*', {'sparse_indices_sharing_key': 'W_h'}),
          ('W_x', {'sparsity_ratio': sparsity_ratio,
                   'initializer': sparse_initializer}),
          ('W_h', {'sparsity_ratio': sparsity_ratio,
                   'initializer': sparse_initializer}),
      ]))
    else:
      if layer_norm:
        tiled_linear_class = tiled_linear.LayerNormedTiledLinear
      else:
        tiled_linear_class = tiled_linear.TiledLinear
      init_params.update(collections.OrderedDict([
          ('W_.*', {'initializer': cell_initializer}),
          ('B_.*', {'initializer': cell_initializer})
      ]))
    def make_layer(layer_index):
      cell = TiledLSTMCell(
          hidden_size_for_layer(layer_index),
          tie_gates=tie_forget_and_input_gates,
          cap_input_gate=cap_input_gate,
          feature_mask_rounds=feature_mask_rounds,
          feature_mask_rank=feature_mask_rank,
          input_transform=input_transforms[layer_index],
          state_transform=state_transforms[layer_index],
          update_transform=update_transforms[layer_index],
          tiled_linear_class=tiled_linear_class,
          tiled_linear_var_init_params=init_params,
          initializer=cell_initializer,
          cell_clip=cell_clip if cell_clip > 0 else None,
          layer_norm=layer_norm,
          activation=eval(activation_fn))  # pylint: disable=eval-used
      return cell
    layers = [make_layer(i) for i in six.moves.range(num_layers)]
    if lstm_skip_connection:
      assert not residual_connections
      return SkipMultiRNNCell(layers)
    elif residual_connections:
      return ResMultiRNNCell(layers)
    else:
      return tf.nn.rnn_cell.MultiRNNCell(layers)

  def make_rhn_column():
    init_params = collections.OrderedDict([
        ('B_c', {'initializer': tf.constant_initializer(forget_bias)}),
    ])

    if overlay_rank > 0:
      assert sparsity_ratio < 0
      # TODO(melisgl): Specify initializers for the shared matrices.
      tiled_linear_class = tiled_linear.OverlayedTiledLinear
      init_params.update(collections.OrderedDict([
          ('W_x_h', {'overlay_sharing_key': 'W_x_any',
                     'overlay_rank': overlay_rank}),
          ('W_x_c', {'overlay_sharing_key': 'W_x_any',
                     'overlay_rank': overlay_rank}),
          ('W_x_t', {'overlay_sharing_key': 'W_x_any',
                     'overlay_rank': overlay_rank}),
          ('W_s_h', {'overlay_sharing_key': 'W_s_any',
                     'overlay_rank': overlay_rank}),
          ('W_s_c', {'overlay_sharing_key': 'W_s_any',
                     'overlay_rank': overlay_rank}),
          ('W_s_t', {'overlay_sharing_key': 'W_s_any',
                     'overlay_rank': overlay_rank}),
      ]))
    elif sparsity_ratio >= 0.0:
      assert overlay_rank == -1
      tiled_linear_class = tiled_linear.SparseTiledLinear
      sparse_initializer = tf.truncated_normal_initializer(
          stddev=math.sqrt(cell_init_factor /
                           sparsity_ratio /
                           # TODO(melisgl): This is off if the input
                           # embedding size is different from the hidden
                           # size.
                           hidden_size))
      init_params.update(collections.OrderedDict([
          ('W_x_.*', {'sparse_indices_sharing_key': 'W_x'}),
          ('W_s_.*', {'sparse_indices_sharing_key': 'W_s'}),
          ('W_x', {'sparsity_ratio': sparsity_ratio,
                   'initializer': sparse_initializer}),
          ('W_s', {'sparsity_ratio': sparsity_ratio,
                   'initializer': sparse_initializer}),
      ]))
    else:
      tiled_linear_class = tiled_linear.TiledLinear
      init_params.update(collections.OrderedDict([
          ('W_.*', {'initializer': cell_initializer}),
      ]))
    logging.info('Creating RHN of depth %s', num_layers)
    if layer_norm:
      logging.warn('RHN does not support layer normalization.')
    cell = TiledRHNCell(
        hidden_size,
        depth=num_layers,
        tie_gates=tie_forget_and_input_gates,
        input_transform=input_transforms[layer_index],
        state_transform=state_transforms[layer_index],
        update_transform=update_transforms[layer_index],
        tiled_linear_class=tiled_linear_class,
        tiled_linear_var_init_params=init_params,
        cell_clip=cell_clip if cell_clip > 0 else None,
        activation=eval(activation_fn))  # pylint: disable=eval-used
    return cell

  def make_nas_column():
    assert not layer_norm
    def make_layer(layer_index):
      logging.info('Creating layer %s', layer_index)

      cell = NASCell(