code: def flatten_batch_to_2d(
      data: type_alias.TensorLike,
      sizes: type_alias.TensorLike = None,
      name: str = "utils_flatten_batch_to_2d",
  ):
      """Reshapes a batch of 2d Tensors by flattening across the batch dimensions.
  
      Note:
        In the following, A1 to An are optional batch dimensions.
  
      A tensor with shape `[A1, ..., An, D1, D2]` will be reshaped to one
      with shape `[A1*...*An*D1, D2]`. This function also returns an inverse
      function that returns any tensor with shape `[A1*...*An*D1, D3]` to one
      with shape `[A1, ..., An, D1, D3]`.
  
      Padded inputs in dimension D1 are allowed. `sizes` determines the first
      elements from D1 to select from each batch dimension.
  
      Examples:
        ```python
        data = [[[1., 2.], [3., 4.]],
                [[5., 6.], [7., 8.]],
                [[9., 10.], [11., 12.]]]
        sizes = None
        output = flatten_batch_to_2d(data, size)
        print(output)
        >>> [[1., 2.], [3., 4.], [5., 6.], [7., 8.], [9., 10.], [11., 12.]]
  
        data = [[[1., 2.], [0., 0.]],
                [[5., 6.], [7., 8.]],
                [[9., 10.], [0., 0.]]]
        sizes = [1, 2, 1]
        output = flatten_batch_to_2d(data, size)
        print(output)
        >>> [[1., 2.], [5., 6.], [7., 8.], [9., 10.]]
        ```
  
      Args:
        data: A tensor with shape `[A1, ..., An, D1, D2]`.
        sizes: An `int` tensor with shape `[A1, ..., An]`. Can be `None`. `sizes[i]
          <= D1`.
        name: A name for this op. Defaults to 'utils_flatten_batch_to_2d'.
  
      Returns:
        A tensor with shape `[A1*...*An*D1, D2]` if `sizes == None`, otherwise a
          tensor  with shape `[sum(sizes), D2]`.
        A function that reshapes a tensor with shape `[A1*...*An*D1, D3]` to a
          tensor with shape `[A1, ..., An, D1, D3]` if `sizes == None`, otherwise
          it reshapes a tensor with shape `[sum(sizes), D3]` to one with shape
          `[A1, ..., An, ..., D1, D3]`.
  
      Raises:
        ValueError: if the input tensor dimensions are invalid.
      """
      with tf.name_scope(name):
          data = tf.convert_to_tensor(value=data)
          if sizes is not None:
              sizes = tf.convert_to_tensor(value=sizes)
  
          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 - 2,
              )
              shape.compare_batch_dimensions(
                  tensors=(data, sizes),
                  tensor_names=("data", "sizes"),
                  last_axes=(-3, -1),
                  broadcast_compatible=False,
              )
  
          data_shape = tf.shape(input=data)
          if sizes is None:
              flat = tf.reshape(data, shape=(-1, data_shape[-1]))
  
              def unflatten(flat, name="utils_unflatten"):
                  """Invert flatten_batch_to_2d."""
                  with tf.name_scope(name):
                      flat = tf.convert_to_tensor(value=flat)
                      output_shape = tf.concat(
                          (data_shape[:-1], tf.shape(input=flat)[-1:]),
                          axis=0,
                      )
                      return tf.reshape(flat, output_shape)
  
          else:
              # Create a mask for the desired rows in `data` to select for flattening:
              # `mask` has shape `[A1, ..., An, D1]` and
              # `mask[a1, ..., an, :] = [True, ..., True, False, ..., False]` where
              # the number of True elements is `sizes[a1, ..., an]`.
              mask = tf.sequence_mask(sizes, data_shape[-2])
              mask_indices = tf.cast(tf.where(mask), tf.int32)
              flat = tf.gather_nd(params=data, indices=mask_indices)
  
              def unflatten(flat, name="utils_unflatten"):
                  """Invert flatten_batch_to_2d."""
                  with tf.name_scope(name):
                      flat = tf.convert_to_tensor(value=flat)
                      output_shape = tf.concat(
                          (data_shape[:-1], tf.shape(input=flat)[-1:]),
                          axis=0,
                      )
                      return tf.scatter_nd(
                          indices=mask_indices,
                          updates=flat,
                          shape=output_shape,
                      )
  
          return flat, unflatten
  
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 5 for calling a global function with example usage: flatten_batch_to_2d(data, sizes, name) and returns: flat,  unflatten, tf, tf