Returns (batched) matmul of a (sparse) matrix with a column vector.
tfp.substrates.jax.math.sparse_or_dense_matvecmul(
sparse_or_dense_matrix,
dense_vector,
validate_args=False,
name=None,
**kwargs
)
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 .
|