def batch_multiply(sp_tensor, dense_vec):
  """Batch multiply a vector with a sparse tensor."""
  batch_indices = sp_tensor.indices[:, 0]
  batch_vec = tf.gather(dense_vec, batch_indices)
  return tf.SparseTensor(
      indices=sp_tensor.indices,
      values=sp_tensor.values * batch_vec,
      dense_shape=sp_tensor.dense_shape)
Batch multiply a vector with a sparse tensor.
language/labs/drkit/model_fns.py