def ragged_sparse_lookup(ragged_ind,
                         ragged_val,
                         row_indices,
                         num_mentions,
                         mask=None,
                         fix_values_to_one=True):
  """Extract rows from sparse matrix and return as another sparse matrix.

  Args:
    ragged_ind: RaggedTensor mapping entities to mention indices. (E x M)
    ragged_val: RaggedTensor mapping entities to mention scores. (E x M)
    row_indices: Batch indices into rows of sparse matrix. (B)
    num_mentions: Scalar (M)
    mask: Boolean mask into the row_indices. Positions where this is false will
      not be retrieved. (B)
    fix_values_to_one: If true final result is binary.

  Returns:
    sp_vec: (B x M) sparse matrix of retrieved rows.
  """
  if mask is None:
    mask = tf.ones_like(row_indices, dtype=tf.bool)

  retained_idx = tf.squeeze(tf.where(mask), 1)
  row_indices = tf.boolean_mask(row_indices, mask)  # NNZ

  batch_idx = tf.gather(ragged_ind, row_indices).to_sparse()  # NNZ
  orig_idx = tf.gather(retained_idx, batch_idx.indices[:, 0])  # B
  batch_idx = tf.concat(
      (tf.expand_dims(orig_idx, 1), tf.expand_dims(batch_idx.values, 1)),
      1)  # B x 2

  if fix_values_to_one:
    return tf.SparseTensor(
        batch_idx,
        tf.ones((tf.shape(batch_idx)[0],), dtype=tf.float32),
        dense_shape=[tf.shape(row_indices)[0], num_mentions])
  else:
    batch_val = tf.gather(ragged_val, row_indices).values  # B
    return tf.SparseTensor(
        batch_idx,
        batch_val,
        dense_shape=[tf.shape(row_indices)[0], num_mentions])
Extract rows from sparse matrix and return as another sparse matrix.

  Args:
    ragged_ind: RaggedTensor mapping entities to mention indices. (E x M)
    ragged_val: RaggedTensor mapping entities to mention scores. (E x M)
    row_indices: Batch indices into rows of sparse matrix. (B)
    num_mentions: Scalar (M)
    mask: Boolean mask into the row_indices. Positions where this is false will
      not be retrieved. (B)
    fix_values_to_one: If true final result is binary.

  Returns:
    sp_vec: (B x M) sparse matrix of retrieved rows.
  
language/labs/drkit/model_fns.py