code: def pooling(
      inputs,
      window_shape,
      pooling_configs=None,
      strides=None,
      padding="VALID",
  ):
      """Applies configurable pooling.
  
      Args:
        inputs: an nd-array; Thego shape of inputs is `[bs, <window dims>,
          features]` and for presence_weights, the shape is `[bs, <window dims>]`.
        window_shape: tuple; Defining the window to reduce over.
        pooling_configs: dict; Configuration for the optional pooling operation.
        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.
      """
      # TODO(dehghani): add positional embedding to other type of pooling?
      strides = strides or window_shape
  
      pooling_type = pooling_configs.get("pooling_type")
      if pooling_type == "avg_pooling":
          x = nn.avg_pool(
              inputs, window_shape, strides=strides, padding=padding
          )
  
      elif pooling_type == "max_pooling":
          x = nn.max_pool(
              inputs, window_shape, strides=strides, padding=padding
          )
  
      elif pooling_type == "space_to_depth":
          x = space_to_depth(
              inputs, window_shape, strides=strides, padding=padding
          )
  
      else:
          raise ValueError(
              "Pooling type {} is not defined.".format(pooling_type)
          )
      return x
  
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 6 for calling a global function with example usage: pooling(inputs, window_shape, pooling_configs, strides, padding) and returns: x