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