def sparse_map(ids, mask, vocab_size):
  """Create mapping from order 3 tensor to vocabulary for copying mechanism.

  This will return a SparseTensor of size (batch_size, vocab_size, num_cells).
  Each column of this sparse tensor will be non-zero at the indices specified by
  the `ids` in that cell.

  For rank 2 tensors expand dims along the axis=2 before passing.

  Args:
    ids: <tf.int32>[batch_size, num_cells, seq_len] Word ids of sequence in each
      cell.
    mask: <tf.float32>[batch_size, num_cells, seq_len] Mask for ids.
    vocab_size: Number of rows in resulting tensor.

  Returns:
    sp_tensor: tf.SparseTensor mapping num_cells to vocab_size.
  """
  batch_size, num_cells, seq_len = (
      tf.shape(ids)[0], tf.shape(ids)[1], tf.shape(ids)[2])

  nonzero = batch_size * num_cells * seq_len

  # Indices.
  ids_x = tf.tile(
      tf.expand_dims(tf.expand_dims(tf.range(batch_size), axis=1), axis=2),
      (1, num_cells, seq_len))
  sp_ids_x = tf.reshape(ids_x, (nonzero, 1))
  sp_ids_y = tf.reshape(ids, (nonzero, 1))
  ids_z = tf.tile(
      tf.expand_dims(tf.expand_dims(tf.range(num_cells), axis=0), axis=2),
      (batch_size, 1, seq_len))
  sp_ids_z = tf.reshape(ids_z, (nonzero, 1))
  sp_ids = tf.concat([sp_ids_x, sp_ids_y, sp_ids_z], axis=1)

  # Values.
  lengths = tf.maximum(
      tf.reduce_sum(mask, axis=2, keepdims=True), 1.)
  values = mask / lengths
  sp_values = tf.reshape(values, (nonzero,))

  # SparseTensor.
  sp_tensor = tf.SparseTensor(
      tf.cast(sp_ids, tf.int64), sp_values,
      tf.cast(tf.stack([batch_size, vocab_size, num_cells]), tf.int64))

  return sp_tensor
Create mapping from order 3 tensor to vocabulary for copying mechanism.

  This will return a SparseTensor of size (batch_size, vocab_size, num_cells).
  Each column of this sparse tensor will be non-zero at the indices specified by
  the `ids` in that cell.

  For rank 2 tensors expand dims along the axis=2 before passing.

  Args:
    ids: <tf.int32>[batch_size, num_cells, seq_len] Word ids of sequence in each
      cell.
    mask: <tf.float32>[batch_size, num_cells, seq_len] Mask for ids.
    vocab_size: Number of rows in resulting tensor.

  Returns:
    sp_tensor: tf.SparseTensor mapping num_cells to vocab_size.
  
language/labs/exemplar_decoding/models/common.py