code: 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),
      )
  
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 17 for calling a global function with example usage: _sparse_tensor_dense_matmul(sp_a, b) and returns: tf