code: def _from_grid(x, patch_sidelength):
      """From grid."""
      (n, gh, gw, fh, fw, _) = x.shape
      fh, fw = patch_sidelength, patch_sidelength
  
      x = x.reshape([n, gh, gw, fh, fw, 3])
      x = x.transpose([0, 1, 3, 2, 4, 5])
      x = x.reshape([n, gh * fh, gw * fw, 3])
      return x
  
imports: import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: _from_grid(x, patch_sidelength) and returns: x