tensorflow_probability/python/math/linalg.py
def _sparse_tensor_dense_matmul(sp_a, b, **kwargs):
  """Returns (batched) matmul of a SparseTensor with a Tensor.

  Args:
    sp_a: `SparseTensor` representing a (batch of) matrices.
    b: `Tensor` representing a (batch of) matrices, with the same batch shape of
      `sp_a`. The shape must be compatible with the shape of `sp_a` and kwargs.
    **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul`.

  Returns:
    product: A dense (batch of) matrix-shaped Tensor of the same batch shape and
    dtype as `sp_a` and `b`. If `sp_a` or `b` is adjointed through `kwargs` then
    the shape is adjusted accordingly.
  """
  batch_shape = _get_shape(sp_a)[:-2]

  # Reshape the SparseTensor into a rank 3 SparseTensors, with the
  # batch shape flattened to a single dimension. If the batch rank is 0, then
  # we add a batch dimension of rank 1.
  sp_a = tf.sparse.reshape(sp_a, tf.concat([[-1], _get_shape(sp_a)[-2:]],
                                           axis=0))
  # Reshape b to stack the batch dimension along the rows.
  b = tf.reshape(b, tf.concat([[-1], _get_shape(b)[-1:]], axis=0))

  # Convert the SparseTensor to a matrix in block diagonal form with blocks of
  # matrices [M, N]. This allow us to use tf.sparse_tensor_dense_matmul which
  # only accepts rank 2 (Sparse)Tensors.
  out = tf.sparse.sparse_dense_matmul(_sparse_block_diag(sp_a), b, **kwargs)

  # Finally retrieve the original batch shape from the resulting rank 2 Tensor.
  # Note that we avoid inferring the final shape from `sp_a` or `b` because we
  # might have transposed one or both of them.
  return tf.reshape(
      out,
      tf.concat([batch_shape, [-1], _get_shape(out)[-1:]], axis=0))
def _sparse_tensor_dense_matmul(sp_a, b, **kwargs):
  """Returns (batched) matmul of a SparseTensor with a Tensor.

  Args:
    sp_a: `SparseTensor` representing a (batch of) matrices.
    b: `Tensor` representing a (batch of) matrices, with the same batch shape of
      `sp_a`. The shape must be compatible with the shape of `sp_a` and kwargs.
    **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul`.

  Returns:
    product: A dense (batch of) matrix-shaped Tensor of the same batch shape and
    dtype as `sp_a` and `b`. If `sp_a` or `b` is adjointed through `kwargs` then
    the shape is adjusted accordingly.
  """
  batch_shape = _get_shape(sp_a)[:-2]

  # Reshape the SparseTensor into a rank 3 SparseTensors, with the
  # batch shape flattened to a single dimension. If the batch rank is 0, then
  # we add a batch dimension of rank 1.
  sp_a = tf.sparse.reshape(sp_a, tf.concat([[-1], _get_shape(sp_a)[-2:]],
                                           axis=0))
  # Reshape b to stack the batch dimension along the rows.
  b = tf.reshape(b, tf.concat([[-1], _get_shape(b)[-1:]], axis=0))

  # Convert the SparseTensor to a matrix in block diagonal form with blocks of
  # matrices [M, N]. This allow us to use tf.sparse_tensor_dense_matmul which
  # only accepts rank 2 (Sparse)Tensors.
  out = tf.sparse.sparse_dense_matmul(_sparse_block_diag(sp_a), b, **kwargs)

  # Finally retrieve the original batch shape from the resulting rank 2 Tensor.
  # Note that we avoid inferring the final shape from `sp_a` or `b` because we
  # might have transposed one or both of them.
  return tf.reshape(
      out,
      tf.concat([batch_shape, [-1], _get_shape(out)[-1:]], axis=0))
Returns (batched) matmul of a SparseTensor with a Tensor.

  Args:
    sp_a: `SparseTensor` representing a (batch of) matrices.
    b: `Tensor` representing a (batch of) matrices, with the same batch shape of
      `sp_a`. The shape must be compatible with the shape of `sp_a` and kwargs.
    **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul`.

  Returns:
    product: A dense (batch of) matrix-shaped Tensor of the same batch shape and
    dtype as `sp_a` and `b`. If `sp_a` or `b` is adjointed through `kwargs` then
    the shape is adjusted accordingly.