code: def _sparse_block_diag(sp_a):
      """Returns a block diagonal rank 2 SparseTensor from a batch of SparseTensors.
  
      Args:
        sp_a: A rank 3 `SparseTensor` representing a batch of matrices.
  
      Returns:
        sp_block_diag_a: matrix-shaped, `float` `SparseTensor` with the same dtype
        as `sparse_or_matrix`, of shape [B * M, B * N] where `sp_a` has shape
        [B, M, N]. Each [M, N] batch of `sp_a` is lined up along the diagonal.
      """
      # Construct the matrix [[M, N], [1, 0], [0, 1]] which would map the index
      # (b, i, j) to (Mb + i, Nb + j). This effectively creates a block-diagonal
      # matrix of dense shape [B * M, B * N].
      # Note that this transformation doesn't increase the number of non-zero
      # entries in the SparseTensor.
      sp_a_shape = tf.convert_to_tensor(_get_shape(sp_a, tf.int64))
      ind_mat = tf.concat(
          [[sp_a_shape[-2:]], tf.eye(2, dtype=tf.int64)], axis=0
      )
      indices = tf.matmul(sp_a.indices, ind_mat)
      dense_shape = sp_a_shape[0] * sp_a_shape[1:]
      return tf.SparseTensor(
          indices=indices, values=sp_a.values, dense_shape=dense_shape
      )
  
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 18 for calling a global function with example usage: _sparse_block_diag(sp_a) and returns: tf