def sparse_ragged_mul(sp_tensor,
                      ragged_ind,
                      ragged_val,
                      batch_size,
                      num_mentions,
                      reduce_fn="sum",
                      threshold=None,
                      fix_values_to_one=True):
  """Multiply two sparse matrices, one of which is in ragged form.

  Args:
    sp_tensor: SparseTensor of shape B x E.
    ragged_ind: RaggedTensor mapping entities to mention indices. (E x M)
    ragged_val: RaggedTensor mapping entities to mention scores. (E x M)
    batch_size: Scalar (B)
    num_mentions: Scalar (M)
    reduce_fn: Aggregation function to use (`max` or `sum`).
    threshold: Scalar. Only entries greater than this will be multiplied from
      the SparseTensor.
    fix_values_to_one: Whether to fix Ragged tensor values to 1.

  Returns:
    sp_vec: (B x M) sparse matrix of retrieved rows.
  """
  row_ind = sp_tensor.indices[:, 0]
  col_ind = sp_tensor.indices[:, 1]
  if threshold is not None:
    mask = tf.greater(sp_tensor.values, threshold)
    row_ind = tf.boolean_mask(row_ind, mask)  # NNZ
    col_ind = tf.boolean_mask(col_ind, mask)  # NNZ
    mask_val = tf.boolean_mask(sp_tensor.values, mask)  # NNZ

  gather_idx = tf.gather(ragged_ind, col_ind).to_sparse()  # NNZ x (M)
  row_gather_idx = tf.gather(row_ind, gather_idx.indices[:, 0])  # NNZ'
  out_sp_indices = tf.concat(
      (tf.expand_dims(row_gather_idx, 1), tf.expand_dims(gather_idx.values, 1)),
      1)

  gather_val = tf.gather(ragged_val, col_ind).to_sparse()  # NNZ x (M)
  row_gather_val = tf.gather(mask_val, gather_val.indices[:, 0])  # NNZ'
  if fix_values_to_one:
    out_sp_values = row_gather_val
  else:
    out_sp_values = row_gather_val * gather_val.values

  agg_indices, agg_values = aggregate_sparse_indices(out_sp_indices,
                                                     out_sp_values,
                                                     [batch_size, num_mentions],
                                                     reduce_fn)

  return tf.SparseTensor(
      agg_indices, agg_values, dense_shape=[batch_size, num_mentions])
Multiply two sparse matrices, one of which is in ragged form.

  Args:
    sp_tensor: SparseTensor of shape B x E.
    ragged_ind: RaggedTensor mapping entities to mention indices. (E x M)
    ragged_val: RaggedTensor mapping entities to mention scores. (E x M)
    batch_size: Scalar (B)
    num_mentions: Scalar (M)
    reduce_fn: Aggregation function to use (`max` or `sum`).
    threshold: Scalar. Only entries greater than this will be multiplied from
      the SparseTensor.
    fix_values_to_one: Whether to fix Ragged tensor values to 1.

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