code: def split_and_pad(
      data_chunks: List[np._typing.NDArray], T: int
  ) -> np._typing.NDArray:
      """Splits and zero-pads data from episodes into a single ndarray with a fixed T-axis.
  
      Processes each data chunk in `data_chunks`, coming from one episode by splitting
      the chunk into smaller sub-chunks, each of a maximum size `T`. If a sub-chunk is
      smaller than `T`, it is right-padded with zeros to match the desired size T.
      All sub-chunks are then re-combined (concatenated) into a single ndarray, which is
      reshaped to include the new time dimension `T` as axis 1 (axis 0 is the batch
      axis). The resulting output array has dimensions (B=number of sub-chunks, T, ...),
      where '...' represents the original dimensions of the input data (excluding the
      batch dimension).
  
      Args:
          data_chunks: A list where each element is a NumPy array representing
              an episode. Each array's shape should be (episode_length, ...)
              where '...' represents any number of additional dimensions.
          T: The desired time dimension size for each chunk.
  
      Returns:
          A np.ndarray containing the reshaped and padded chunks. The shape of the
          array will be (B, T, ...) where B is automatically determined by the number
          of chunks in `data_chunks` and `T`.
          '...' represents the original dimensions of the input data, excluding the
          batch dimension.
      """
      all_chunks = []
  
      for data_chunk in data_chunks:
          num_sub_chunks = int(np.ceil(data_chunk.shape[0] / T))
  
          for i in range(num_sub_chunks):
              start_index = i * T
              end_index = start_index + T
  
              # Extract the chunk.
              sub_chunk = data_chunk[start_index:end_index]
  
              # Pad the chunk if it's shorter than T
              if sub_chunk.shape[0] < T:
                  padding_shape = [(0, T - sub_chunk.shape[0])] + [
                      (0, 0) for _ in range(sub_chunk.ndim - 1)
                  ]
                  sub_chunk = np.pad(
                      sub_chunk,
                      pad_width=padding_shape,
                      mode="constant",
                  )
  
              all_chunks.append(sub_chunk)
  
      # Combine all chunks into a single array.
      result = np.concatenate(all_chunks, axis=0)
  
      # Reshape the array to include the time dimension T.
      # The new shape should be (-1, T) + original dimensions (excluding the
      # batch dimension).
      result = result.reshape((-1, T) + result.shape[1:])
  
      return result
  
imports: from typing import List, Tuple
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: split_and_pad(data_chunks, T) and returns: result