code: def space_to_depth(
      inputs, window_shape, strides=None, padding="VALID"
  ):
      """Applies space to depth.
  
      Args:
        inputs: Input data with dimensions `[bs, window dims, ..., features]`.
        window_shape: tuple; Defining the window to reduce over.
        strides: tuple, A sequence of `n` integers, representing the inter-window
            strides (default: window_shape).
        padding: str; Either `'SAME'`, `'VALID'`, or a sequence of `n` `(low,
          high)` integer pairs that give the padding to
          apply before and after each spatial dimension (default: `'VALID'`).
  
      Returns:
        An output image with less or equal spacial dimensions as inputs.
  
      """
      strides = strides or window_shape
      patched = extract_image_patches(
          lhs=inputs.astype(jnp.float64),
          rhs_shape=(1,) + window_shape + (1,),
          strides=(1,) + strides + (1,),
          padding=padding,
          rhs_dilation=(1,) * inputs.ndim,
          data_format="NHWC",
      )
  
      bs, n_patch_h, n_patch_w, _, _, _ = patched.shape
      return patched.reshape(bs, n_patch_h, n_patch_w, -1)
  
file_path: scenic/scenic/model_lib/layers/nn_ops.py
imports: from typing import Callable, Sequence
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 5 for calling a global function with example usage: space_to_depth(inputs, window_shape, strides, padding) and returns: patched