def flatten(t):
  """Collapse all dimensions of a tensor but the last.

  This function also returns a function to unflatten any tensors to recover
  the original shape (aside from the last dimension). This is useful for
  interfacing with functions that are agnostic to all dimensions but the last.
  For example, if we want to apply a linear projection to a batched sequence
  of embeddings:

    t = tf.random_uniform([batch_size, sequence_length, embedding_size])
    w = tf.get_variable("w", [embedding_size, projection_size])
    flat_t, unflatten = flatten(t)
    flat_projected = tf.matmul(flat_t, w)
    projected = unflatten(flat_projected)

  Args:
    t: [dim_1, ..., dim_(n-1), dim_n]

  Returns:
    output: [dim_1 * ... * dim_(n-1), dim_n]
    _unflatten: A function that when called with a flattened tensor returns the
        unflattened version, i.e. reshapes any tensor with shape
        [dim_1 * ... * dim_(n-1), dim_new] to [dim_1, ..., dim_(n-1), dim_new].
  """
  input_shape = shape(t)
  if len(input_shape) > 2:
    t = tf.reshape(t, [-1, input_shape[-1]])

  def _unflatten(flat_t):
    if len(input_shape) > 2:
      return tf.reshape(flat_t, input_shape[:-1] + [shape(flat_t, -1)])
    else:
      return flat_t

  return t, _unflatten
Collapse all dimensions of a tensor but the last.

  This function also returns a function to unflatten any tensors to recover
  the original shape (aside from the last dimension). This is useful for
  interfacing with functions that are agnostic to all dimensions but the last.
  For example, if we want to apply a linear projection to a batched sequence
  of embeddings:

    t = tf.random_uniform([batch_size, sequence_length, embedding_size])
    w = tf.get_variable("w", [embedding_size, projection_size])
    flat_t, unflatten = flatten(t)
    flat_projected = tf.matmul(flat_t, w)
    projected = unflatten(flat_projected)

  Args:
    t: [dim_1, ..., dim_(n-1), dim_n]

  Returns:
    output: [dim_1 * ... * dim_(n-1), dim_n]
    _unflatten: A function that when called with a flattened tensor returns the
        unflattened version, i.e. reshapes any tensor with shape
        [dim_1 * ... * dim_(n-1), dim_new] to [dim_1, ..., dim_(n-1), dim_new].
  
language/common/utils/tensor_utils.py