code: def space_to_channel2d(
      x: JaxArray, size: Union[Tuple[int, int], int] = 2
  ) -> JaxArray:
      """Transfer spatial dimensions (H, W) into channel dimension C.
  
      Args:
          x: input tensor of shape (N, C, H, W).
          size: size of spatial area.
  
      Returns:
          output tensor of shape (N, C * size[0] * size[1]), H // size[0], W // size[1]).
      """
      size = to_tuple(size, 2)
      s = x.shape
      y = x.reshape(
          (
              s[0],
              s[1],
              s[2] // size[0],
              size[0],
              s[3] // size[1],
              size[1],
          )
      )
      y = y.transpose((0, 1, 3, 5, 2, 4))
      return y.reshape(
          (
              s[0],
              s[1] * size[0] * size[1],
              s[2] // size[0],
              s[3] // size[1],
          )
      )
  
file_path: objax/objax/functional/core/pooling.py
imports: from typing import Union, Tuple, Optional
  import numpy
  from jax import numpy, lax
  from objax.constants import ConvPadding
  from objax.typing import JaxArray, ConvPaddingInt
  from objax.util import to_padding, to_tuple
label_desc: Use image node_id 6 for calling a global function with example usage: space_to_channel2d(x, size) and returns: y