code: def _to_grid(x, patch_sidelength):
      """To grid."""
      n, h, w, three = x.shape
      assert three == 3
      gh = h // patch_sidelength
      gw = w // patch_sidelength
      fh, fw = patch_sidelength, patch_sidelength
  
      x = jnp.reshape(x, [n, gh, fh, gw, fw, 3])
      x = jnp.transpose(x, [0, 1, 3, 2, 4, 5])
      x_grid = jnp.reshape(x, [n, gh, gw, fh, fw, 3])
      return x_grid
  
imports: import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: _to_grid(x, patch_sidelength) and returns: x_grid