code: def unflatten_2d_to_batch(
      data: type_alias.TensorLike,
      sizes: type_alias.TensorLike,
      max_rows: Optional[int] = None,
      name: str = "utils_unflatten_2d_to_batch",
  ):
      r"""Reshapes a 2d Tensor into a batch of 2d Tensors.
  
      The `data` tensor with shape `[D1, D2]` will be mapped to a tensor with shape
      `[A1, ..., An, max_rows, D2]` where `max_rows` defaults to `max(sizes)`.
      `sizes` determines the segment of rows in the input that get mapped to a
      particular batch dimension (`sum(sizes) == D1`).
  
      Examples:
  
        ```python
        data = [[1., 2.],
                [3., 4.],
                [5., 6.],
                [7., 8.],
                [9., 10.],
                [11., 12.]]
        sizes = [2, 3, 1]
  
        output = unflatten_2d_to_batch(data, sizes, max_rows=None)
        print(output.shape)
        >>> [3, 3, 2]
        print(output)
        >>> [[[1., 2.],
              [3., 4.],
              [0., 0.]],
             [[5., 6.],
              [7., 8.],
              [9., 10.]],
             [[11., 12.],
              [0., 0.],
              [0., 0.]]]
  
        output = unflatten_2d_to_batch(data, sizes, max_rows=4)
        print(output.shape)
        >>> [3, 4, 2]
        print(output)
        >>> [[[1., 2.],
              [3., 4.],
              [0., 0.],
              [0., 0.]],
             [[5., 6.],
              [7., 8.],
              [9., 10.],
              [0., 0.]],
             [[11., 12.],
              [0., 0.],
              [0., 0.],
              [0., 0.]]]
        ```
  
      Args:
        data: A tensor with shape `[D1, D2]`.
        sizes: An `int` tensor with shape `[A1, ..., An]`.
        max_rows: An `int` specifying the maximum number of rows in the unflattened
          output. `max_rows >= max(sizes)`.
        name: A name for this op. Defaults to 'utils_unflatten_2d_to_batch'.
  
      Returns:
        A tensor with shape `[A1, A2, ..., max_rows, D2]`.
      """
      with tf.name_scope(name):
          data = tf.convert_to_tensor(value=data)
          sizes = tf.convert_to_tensor(value=sizes)
          if max_rows is None:
              max_rows = tf.reduce_max(input_tensor=sizes)
          else:
              max_rows = tf.convert_to_tensor(value=max_rows)
  
          shape.check_static(
              tensor=data, tensor_name="data", has_rank=2
          )
          if not sizes.dtype.is_integer:
              raise TypeError("'sizes' must have an integer type.")
  
          mask = tf.sequence_mask(sizes, max_rows)
          mask_indices = tf.cast(tf.where(mask), tf.int32)
          output_shape = tf.concat(
              (
                  tf.shape(input=sizes),
                  (max_rows,),
                  tf.shape(input=data)[-1:],
              ),
              axis=0,
          )
          return tf.scatter_nd(
              indices=mask_indices, updates=data, shape=output_shape
          )
  
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 6 for calling a global function with example usage: unflatten_2d_to_batch(data, sizes, max_rows, name) and returns: tf