code: def convert_to_block_diag_2d(
      data: tf.sparse.SparseTensor,
      sizes: Optional[type_alias.TensorLike] = None,
      validate_indices: bool = False,
      name: str = "utils_convert_to_block_diag_2d",
  ):
      """Convert a batch of 2d SparseTensors to a 2d block diagonal SparseTensor.
  
      Note:
        In the following, A1 to An are optional batch dimensions.
  
      A `SparseTensor` with dense shape `[A1, ..., An, D1, D2]` will be reshaped
      to one with shape `[A1*...*An*D1, A1*...*An*D2]`.
  
      Padded inputs in dims D1 and D2 are allowed. `sizes` indicates the un-padded
      shape for each inner `[D1, D2]` matrix. The additional (padded) rows and
      columns will be omitted in the block diagonal output.
  
      If padded (`sizes != None`), the input should not contain any sparse indices
      outside the bounds indicated by `sizes`. Setting `validate_indices=True` will
      explicitly filter any invalid sparse indices before block diagonalization.
  
      Args:
        data: A `SparseTensor` with dense shape `[A1, ..., An, D1, D2]`.
        sizes: A tensor with shape `[A1, ..., An, 2]`. Can be `None` (indicates no
          padding). If not `None`, `sizes` indicates the true sizes (before padding)
          of the inner dimensions of `data`.
        validate_indices: A boolean. Ignored if `sizes==None`. If True,
          out-of-bounds indices in `data` are explicitly ignored, otherwise
          out-of-bounds indices will cause undefined behavior.
        name: A name for this op. Defaults to 'utils_convert_to_block_diag_2d'.
  
      Returns:
        A 2d block-diagonal SparseTensor.
  
      Raises:
        TypeError: if the input types are invalid.
        ValueError: if the input dimensions are invalid.
      """
      with tf.name_scope(name):
          data = tf.compat.v1.convert_to_tensor_or_sparse_tensor(
              value=data
          )
          if sizes is not None:
              sizes = tf.convert_to_tensor(value=sizes)
  
          if not isinstance(data, tf.SparseTensor):
              raise TypeError("'data' must be a 'SparseTensor'.")
          if sizes is not None and not sizes.dtype.is_integer:
              raise TypeError("'sizes' must have an integer type.")
          shape.check_static(
              tensor=data, tensor_name="data", has_rank_greater_than=2
          )
          if sizes is not None:
              shape.check_static(
                  tensor=sizes,
                  tensor_name="sizes",
                  has_rank=data.shape.ndims - 1,
                  has_dim_equals=(-1, 2),
              )
              shape.compare_batch_dimensions(
                  tensors=(data, sizes),
                  tensor_names=("data", "sizes"),
                  last_axes=(-3, -2),
                  broadcast_compatible=False,
              )
  
          data_shape = tf.shape(input=data)
          data = tf.sparse.reshape(
              data, [-1, data_shape[-2], data_shape[-1]]
          )
          indices = data.indices
          if sizes is not None:
              sizes = tf.cast(
                  tf.reshape(sizes, shape=(-1, 2)), tf.int64
              )
              if validate_indices:
                  in_bounds = ~tf.reduce_any(
                      input_tensor=indices[:, 1:]
                      >= tf.gather(sizes, indices[:, 0]),
                      axis=-1,
                  )
                  indices = tf.boolean_mask(
                      tensor=indices, mask=in_bounds
                  )
                  values = tf.boolean_mask(
                      tensor=data.values, mask=in_bounds
                  )
              else:
                  values = data.values
              cumsum = tf.cumsum(sizes, axis=0, exclusive=True)
              index_shift = tf.gather(cumsum, indices[:, 0])
              indices = indices[:, 1:] + index_shift
              block_diag = tf.SparseTensor(
                  indices,
                  values,
                  tf.reduce_sum(input_tensor=sizes, axis=0),
              )
          else:
              data_shape = tf.shape(input=data, out_type=tf.int64)
              index_shift = (
                  tf.expand_dims(indices[:, 0], -1) * data_shape[1:]
              )
              indices = indices[:, 1:] + index_shift
              block_diag = tf.SparseTensor(
                  indices, data.values, data_shape[0] * data_shape[1:]
              )
          return block_diag
  
file_path: graphics/tensorflow_graphics/geometry/convolution/utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Any, List, Optional, Tuple, Union
  import tensorflow
  from tensorflow_graphics.util import shape
  from tensorflow_graphics.util import type_alias
label_desc: Use image node_id 7 for calling a global function with example usage: convert_to_block_diag_2d(data, sizes, validate_indices, name) and returns: block_diag