code: def _normalize_by_window_size(
      dims, spatial_strides, padding
  ):  # pylint: disable=invalid-name
      def rescale(outputs, inputs):
          one = jnp.ones(inputs.shape[1:-1], dtype=inputs.dtype)
          window_sizes = lax.reduce_window(
              one, 0.0, lax.add, dims, spatial_strides, padding
          )
          return outputs / window_sizes[..., jnp.newaxis]
  
      return rescale
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import jax.scipy.special
  import numpy
  import tensorflow
  import tensorflow_datasets
  from trax.fastmath import numpy
  from trax.shapes import signature
label_desc: Use image node_id 5 for calling a global function with example usage: _normalize_by_window_size(dims, spatial_strides, padding) and returns: rescale, unknown