def _prepare_memory(
    memory,
    memory_sequence_length,
    mask,
    check_inner_dims_defined):
  """Convert to tensor and possibly mask `memory`.

  Args:
    memory: `Tensor`, shaped `[batch_size, max_time, ...]`.
    memory_sequence_length: `int32` `Tensor`, shaped `[batch_size]`.
    mask: To mask out some of the elements.
    check_inner_dims_defined: Python boolean.  If `True`, the `memory`
      argument's shape is checked to ensure all but the two outermost
      dimensions are fully defined.

  Returns:
    A (possibly masked), checked, new `memory`.

  Raises:
    ValueError: If `check_inner_dims_defined` is `True` and not
      `memory.shape[2:].is_fully_defined()`.
  """
  memory = tf.contrib.framework.nest.map_structure(
      lambda m: tf.convert_to_tensor(m, name="memory"), memory)
  if memory_sequence_length is not None:
    memory_sequence_length = tf.convert_to_tensor(
        memory_sequence_length, name="memory_sequence_length")
  if check_inner_dims_defined:
    def _check_dims(m):
      if not m.get_shape()[2:].is_fully_defined():
        raise ValueError("Expected memory %s to have fully defined inner dims, "
                         "but saw shape: %s" % (m.name, m.get_shape()))

    tf.contrib.framework.nest.map_structure(_check_dims, memory)

  seq_len_mask = tf.cast(mask,
                         tf.contrib.framework.nest.flatten(memory)[0].dtype)
  seq_len_batch_size = (dimension_value(mask.shape[0]) or tf.shape(mask)[0])

  def _maybe_mask(m, seq_len_mask):
    """Mask the sequence with m."""
    rank = m.get_shape().ndims
    rank = rank if rank is not None else tf.rank(m)
    extra_ones = tf.ones(rank - 2, dtype=tf.int32)
    m_batch_size = dimension_value(m.shape[0]) or tf.shape(m)[0]
    with tf.control_dependencies(
        [tf.assert_equal(seq_len_batch_size, m_batch_size, message="batch")]):
      seq_len_mask = tf.reshape(
          seq_len_mask, tf.concat((tf.shape(seq_len_mask), extra_ones), 0))
      return m * seq_len_mask

  return tf.contrib.framework.nest.map_structure(
      lambda m: _maybe_mask(m, seq_len_mask), memory)
Convert to tensor and possibly mask `memory`.

  Args:
    memory: `Tensor`, shaped `[batch_size, max_time, ...]`.
    memory_sequence_length: `int32` `Tensor`, shaped `[batch_size]`.
    mask: To mask out some of the elements.
    check_inner_dims_defined: Python boolean.  If `True`, the `memory`
      argument's shape is checked to ensure all but the two outermost
      dimensions are fully defined.

  Returns:
    A (possibly masked), checked, new `memory`.

  Raises:
    ValueError: If `check_inner_dims_defined` is `True` and not
      `memory.shape[2:].is_fully_defined()`.
  
language/labs/exemplar_decoding/models/attention.py