tensorflow_probability/python/math/linalg.py
def sparse_or_dense_matvecmul(sparse_or_dense_matrix,
                              dense_vector,
                              validate_args=False,
                              name=None,
                              **kwargs):
  """Returns (batched) matmul of a (sparse) matrix with a column vector.

  Args:
    sparse_or_dense_matrix: `SparseTensor` or `Tensor` representing a (batch of)
      matrices.
    dense_vector: `Tensor` representing a (batch of) vectors, with the same
      batch shape as `sparse_or_dense_matrix`. The shape must be compatible with
      the shape of `sparse_or_dense_matrix` and kwargs.
    validate_args: When `True`, additional assertions might be embedded in the
      graph.
      Default value: `False` (i.e., no graph assertions are added).
    name: Python `str` prefixed to ops created by this function.
      Default value: "sparse_or_dense_matvecmul".
    **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul` or
      `tf.matmul`.

  Returns:
    product: A dense (batch of) vector-shaped Tensor of the same batch shape and
    dtype as `sparse_or_dense_matrix` and `dense_vector`.
  """
  with tf.compat.v1.name_scope(name, 'sparse_or_dense_matvecmul',
                               [sparse_or_dense_matrix, dense_vector]):
    dense_vector = tf.convert_to_tensor(
        value=dense_vector, dtype_hint=tf.float32, name='dense_vector')
    return tf.squeeze(
        sparse_or_dense_matmul(
            sparse_or_dense_matrix,
            dense_vector[..., tf.newaxis],
            validate_args=validate_args,
            **kwargs),
        axis=[-1])
def sparse_or_dense_matvecmul(sparse_or_dense_matrix,
                              dense_vector,
                              validate_args=False,
                              name=None,
                              **kwargs):
  """Returns (batched) matmul of a (sparse) matrix with a column vector.

  Args:
    sparse_or_dense_matrix: `SparseTensor` or `Tensor` representing a (batch of)
      matrices.
    dense_vector: `Tensor` representing a (batch of) vectors, with the same
      batch shape as `sparse_or_dense_matrix`. The shape must be compatible with
      the shape of `sparse_or_dense_matrix` and kwargs.
    validate_args: When `True`, additional assertions might be embedded in the
      graph.
      Default value: `False` (i.e., no graph assertions are added).
    name: Python `str` prefixed to ops created by this function.
      Default value: "sparse_or_dense_matvecmul".
    **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul` or
      `tf.matmul`.

  Returns:
    product: A dense (batch of) vector-shaped Tensor of the same batch shape and
    dtype as `sparse_or_dense_matrix` and `dense_vector`.
  """
  with tf.compat.v1.name_scope(name, 'sparse_or_dense_matvecmul',
                               [sparse_or_dense_matrix, dense_vector]):
    dense_vector = tf.convert_to_tensor(
        value=dense_vector, dtype_hint=tf.float32, name='dense_vector')
    return tf.squeeze(
        sparse_or_dense_matmul(
            sparse_or_dense_matrix,
            dense_vector[..., tf.newaxis],
            validate_args=validate_args,
            **kwargs),
        axis=[-1])
Returns (batched) matmul of a (sparse) matrix with a column vector.

  Args:
    sparse_or_dense_matrix: `SparseTensor` or `Tensor` representing a (batch of)
      matrices.
    dense_vector: `Tensor` representing a (batch of) vectors, with the same
      batch shape as `sparse_or_dense_matrix`. The shape must be compatible with
      the shape of `sparse_or_dense_matrix` and kwargs.
    validate_args: When `True`, additional assertions might be embedded in the
      graph.
      Default value: `False` (i.e., no graph assertions are added).
    name: Python `str` prefixed to ops created by this function.
      Default value: "sparse_or_dense_matvecmul".
    **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul` or
      `tf.matmul`.

  Returns:
    product: A dense (batch of) vector-shaped Tensor of the same batch shape and
    dtype as `sparse_or_dense_matrix` and `dense_vector`.