code: def patched_color_jitter(
      input_image, aug_params, aug_fn=fast_color_perturb
  ):
      """Color jitter applied to patch granularity."""
      (
          local_batch_size,
          num_patches,
          num_patches_1,
          three,
      ) = aug_params.shape
      assert three == 3
      assert num_patches == num_patches_1
  
      local_batch_size_1, h, w, three = input_image.shape
      assert three == 3
      assert h == w
      # assert h == 224
      assert local_batch_size_1 == local_batch_size
  
      patch_sidelength = h // num_patches
      grid = _to_grid(input_image, patch_sidelength)
  
      fast_color_jitter_vmapped = jax.vmap(
          aug_fn, in_axes=0, out_axes=0
      )
      fast_color_jitter_vmapped = jax.vmap(
          fast_color_jitter_vmapped, in_axes=1, out_axes=1
      )
      fast_color_jitter_vmapped = jax.vmap(
          fast_color_jitter_vmapped, in_axes=2, out_axes=2
      )
  
      jittered_grid = fast_color_jitter_vmapped(grid, aug_params)
      jittered = _from_grid(jittered_grid, patch_sidelength)
      return jittered
  
imports: import jax
  import jax.numpy
label_desc: Use image node_id 4 for calling a global function with example usage: patched_color_jitter(input_image, aug_params, aug_fn) and returns: jittered