code: def max_pool(
      value: jax.Array,
      window_shape: Union[int, Sequence[int]],
      strides: Union[int, Sequence[int]],
      padding: str,
      channel_axis: Optional[int] = -1,
  ) -> jax.Array:
      """Max 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.
      """
      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)
  
      return lax.reduce_window(
          value, -jnp.inf, lax.max, window_shape, strides, padding
      )
  
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 3 for calling a global function with example usage: max_pool(value, window_shape, strides, padding, channel_axis) and returns: lax