tensorflow_probability/python/math/linalg.py
def sparse_or_dense_matmul(sparse_or_dense_a,
                           dense_b,
                           validate_args=False,
                           name=None,
                           **kwargs):
  """Returns (batched) matmul of a SparseTensor (or Tensor) with a Tensor.

  Args:
    sparse_or_dense_a: `SparseTensor` or `Tensor` representing a (batch of)
      matrices.
    dense_b: `Tensor` representing a (batch of) matrices, with the same batch
      shape as `sparse_or_dense_a`. The shape must be compatible with the shape
      of `sparse_or_dense_a` 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_matmul".
    **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul` or
      `tf.matmul`.

  Returns:
    product: A dense (batch of) matrix-shaped Tensor of the same batch shape and
    dtype as `sparse_or_dense_a` and `dense_b`. If `sparse_or_dense_a` or
    `dense_b` is adjointed through `kwargs` then the shape is adjusted
    accordingly.
  """
  with tf.compat.v1.name_scope(name, 'sparse_or_dense_matmul',
                               [sparse_or_dense_a, dense_b]):
    dense_b = tf.convert_to_tensor(
        value=dense_b, dtype_hint=tf.float32, name='dense_b')

    if validate_args:
      assert_a_rank_at_least_2 = tf.compat.v1.assert_rank_at_least(
          sparse_or_dense_a,
          rank=2,
          message='Input `sparse_or_dense_a` must have at least 2 dimensions.')
      assert_b_rank_at_least_2 = tf.compat.v1.assert_rank_at_least(
          dense_b,
          rank=2,
          message='Input `dense_b` must have at least 2 dimensions.')
      with tf.control_dependencies(
          [assert_a_rank_at_least_2, assert_b_rank_at_least_2]):
        sparse_or_dense_a = tf.identity(sparse_or_dense_a)
        dense_b = tf.identity(dense_b)

    if isinstance(sparse_or_dense_a,
                  (tf.SparseTensor, tf.compat.v1.SparseTensorValue)):
      return _sparse_tensor_dense_matmul(sparse_or_dense_a, dense_b, **kwargs)
    else:
      return tf.matmul(sparse_or_dense_a, dense_b, **kwargs)
def sparse_or_dense_matmul(sparse_or_dense_a,
                           dense_b,
                           validate_args=False,
                           name=None,
                           **kwargs):
  """Returns (batched) matmul of a SparseTensor (or Tensor) with a Tensor.

  Args:
    sparse_or_dense_a: `SparseTensor` or `Tensor` representing a (batch of)
      matrices.
    dense_b: `Tensor` representing a (batch of) matrices, with the same batch
      shape as `sparse_or_dense_a`. The shape must be compatible with the shape
      of `sparse_or_dense_a` 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_matmul".
    **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul` or
      `tf.matmul`.

  Returns:
    product: A dense (batch of) matrix-shaped Tensor of the same batch shape and
    dtype as `sparse_or_dense_a` and `dense_b`. If `sparse_or_dense_a` or
    `dense_b` is adjointed through `kwargs` then the shape is adjusted
    accordingly.
  """
  with tf.compat.v1.name_scope(name, 'sparse_or_dense_matmul',
                               [sparse_or_dense_a, dense_b]):
    dense_b = tf.convert_to_tensor(
        value=dense_b, dtype_hint=tf.float32, name='dense_b')

    if validate_args:
      assert_a_rank_at_least_2 = tf.compat.v1.assert_rank_at_least(
          sparse_or_dense_a,
          rank=2,
          message='Input `sparse_or_dense_a` must have at least 2 dimensions.')
      assert_b_rank_at_least_2 = tf.compat.v1.assert_rank_at_least(
          dense_b,
          rank=2,
          message='Input `dense_b` must have at least 2 dimensions.')
      with tf.control_dependencies(
          [assert_a_rank_at_least_2, assert_b_rank_at_least_2]):
        sparse_or_dense_a = tf.identity(sparse_or_dense_a)
        dense_b = tf.identity(dense_b)

    if isinstance(sparse_or_dense_a,
                  (tf.SparseTensor, tf.compat.v1.SparseTensorValue)):
      return _sparse_tensor_dense_matmul(sparse_or_dense_a, dense_b, **kwargs)
    else:
      return tf.matmul(sparse_or_dense_a, dense_b, **kwargs)
def sparse_or_dense_matmul ( sparse_or_dense_a , dense_b , validate_args = False , name = None , * * kwargs ) : with tf . compat . v1 . name_scope ( name , 'sparse_or_dense_matmul' , [ sparse_or_dense_a , dense_b ] ) : dense_b = tf . convert_to_tensor ( value = dense_b , dtype_hint = tf . float32 , name = 'dense_b' ) if validate_args : assert_a_rank_at_least_2 = tf . compat . v1 . assert_rank_at_least ( sparse_or_dense_a , rank = 2 , message = 'Input `sparse_or_dense_a` must have at least 2 dimensions.' ) assert_b_rank_at_least_2 = tf . compat . v1 . assert_rank_at_least ( dense_b , rank = 2 , message = 'Input `dense_b` must have at least 2 dimensions.' ) with tf . control_dependencies ( [ assert_a_rank_at_least_2 , assert_b_rank_at_least_2 ] ) : sparse_or_dense_a = tf . identity ( sparse_or_dense_a ) dense_b = tf . identity ( dense_b ) if isinstance ( sparse_or_dense_a , ( tf . SparseTensor , tf . compat . v1 . SparseTensorValue ) ) : return _sparse_tensor_dense_matmul ( sparse_or_dense_a , dense_b , * * kwargs ) else : return tf . matmul ( sparse_or_dense_a , dense_b , * * kwargs )
Returns (batched) matmul of a SparseTensor (or Tensor) with a Tensor.

  Args:
    sparse_or_dense_a: `SparseTensor` or `Tensor` representing a (batch of)
      matrices.
    dense_b: `Tensor` representing a (batch of) matrices, with the same batch
      shape as `sparse_or_dense_a`. The shape must be compatible with the shape
      of `sparse_or_dense_a` 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_matmul".
    **kwargs: Keyword arguments to `tf.sparse_tensor_dense_matmul` or
      `tf.matmul`.

  Returns:
    product: A dense (batch of) matrix-shaped Tensor of the same batch shape and
    dtype as `sparse_or_dense_a` and `dense_b`. If `sparse_or_dense_a` or
    `dense_b` is adjointed through `kwargs` then the shape is adjusted
    accordingly.
Returns ( batched ) matmul of a SparseTensor ( or Tensor ) with a Tensor .