def sparse_batched_matmul(x, y):
  """Batch multiply sparse tensor x with dense tensor y.

  Args:
    x: <tf.float32>[B, M, N] SparseTensor.
    y: <tf.float32>[B, N, K] DenseTensor.

  Returns:
    <tf.float32>[B, M, K] DenseTensor.
  """
  sp_indices = x.indices
  sp_values = x.values
  batch_size, num_row = x.dense_shape[0], x.dense_shape[1]
  num_col = tf.cast(tf.shape(y)[2], tf.int64)
  # Number of non-zero entries.
  num_nz = tf.cast(tf.shape(sp_indices)[0], tf.int64)

  # Fetch relevant values from y.
  # <tf.float32>[num_nz, num_col]
  lookup = tf.gather_nd(y, tf.stack(
      [sp_indices[:, 0], sp_indices[:, 2]], axis=1))

  # Reshape first two dimensions of x into a new SparseTensor of rank 2.
  # <tf.int32>[num_nz, 2]
  x_i = tf.stack([sp_indices[:, 0] * num_row + sp_indices[:, 1],
                  tf.cast(tf.range(num_nz), tf.int64)], axis=1)
  # <tf.float32>[batch_size * num_row, num_nz]
  sparse_2d = tf.SparseTensor(x_i, sp_values,
                              tf.stack([batch_size * num_row, num_nz], 0))

  # Multiply the new sparse tensor with the gathered values.
  # <tf.float32>[batch_size * num_row, num_col]
  dense_2d = tf.sparse_tensor_dense_matmul(sparse_2d, lookup)

  # Reshape back [batch_size, num_row, num_col]
  dense_3d = tf.reshape(dense_2d, tf.stack([batch_size, num_row, num_col], 0))

  return dense_3d
Batch multiply sparse tensor x with dense tensor y.

  Args:
    x: <tf.float32>[B, M, N] SparseTensor.
    y: <tf.float32>[B, N, K] DenseTensor.

  Returns:
    <tf.float32>[B, M, K] DenseTensor.
  
language/labs/exemplar_decoding/models/common.py