code: 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.name_scope(name or "sparse_or_dense_matmul"):
          dense_b = tf.convert_to_tensor(
              dense_b, dtype_hint=tf.float32, name="dense_b"
          )
  
          if validate_args:
              assert_a_rank_at_least_2 = assert_util.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 = assert_util.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, tf1.SparseTensorValue),
          ):
              return _sparse_tensor_dense_matmul(
                  sparse_or_dense_a, dense_b, **kwargs
              )
          else:
              return tf.matmul(sparse_or_dense_a, dense_b, **kwargs)
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
label_desc: Use image node_id 12 for calling a global function with example usage: sparse_or_dense_matmul(sparse_or_dense_a, dense_b, validate_args, name) and returns: _sparse_tensor_dense_matmul, tf