def _sparse_block_diag(sp_a):
<EOL>sp_a_shape = tf.convert_to_tensor(value=_get_shape(sp_a, tf.int64))<EOL>ind_mat = tf.concat([[sp_a_shape[-<NUM_LIT:2>:]], tf.eye(<NUM_LIT:2>, dtype=tf.int64)], axis=<NUM_LIT:0>)<EOL>indices = tf.matmul(sp_a.indices, ind_mat)<EOL>dense_shape = sp_a_shape[<NUM_LIT:0>] * sp_a_shape[<NUM_LIT:1>:]<EOL>return tf.SparseTensor(<EOL>indices=indices, values=sp_a.values, dense_shape=dense_shape)<EOL>
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.