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`.

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))