code: def avg_pool(
      value: jax.Array,
      window_shape: Union[int, Sequence[int]],
      strides: Union[int, Sequence[int]],
      padding: str,
      channel_axis: Optional[int] = -1,
  ) -> jax.Array:
      """Average pool.
  
      Args:
        value: Value to pool.
        window_shape: Shape of the pooling window, same rank as value.
        strides: Strides of the pooling window, same rank as value.
        padding: Padding algorithm. Either ``VALID`` or ``SAME``.
        channel_axis: Axis of the spatial channels for which pooling is skipped.
  
      Returns:
        Pooled result. Same rank as value.
  
      Raises:
        ValueError: If the padding is not valid.
      """
      if padding not in ("SAME", "VALID"):
          raise ValueError(
              f"Invalid padding '{padding}', must be 'SAME' or 'VALID'."
          )
  
      _warn_if_unsafe(window_shape, strides)
      window_shape = _infer_shape(value, window_shape, channel_axis)
      strides = _infer_shape(value, strides, channel_axis)
  
      reduce_window_args = (
          0.0,
          lax.add,
          window_shape,
          strides,
          padding,
      )
      pooled = lax.reduce_window(value, *reduce_window_args)
      if padding == "VALID":
          # Avoid the extra reduce_window.
          return pooled / np.prod(window_shape)
      else:
          # Count the number of valid entries at each input point, then use that for
          # computing average. Assumes that any two arrays of same shape will be
          # padded the same. Avoid broadcasting on axis where pooling is skipped.
          shape = [
              (v if w != 1 else 1)
              for (v, w) in zip(value.shape, window_shape)
          ]
          window_counts = lax.reduce_window(
              jnp.ones(shape, value.dtype), *reduce_window_args
          )
          return pooled / window_counts
  
imports: from collections.abc import Sequence
  from typing import Optional, Union
  import warnings
  from haiku._src import module
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: avg_pool(value, window_shape, strides, padding, channel_axis) and returns: unknown, unknown