code: def max_pool_2d(
      x: JaxArray,
      size: Union[Tuple[int, int], int] = 2,
      strides: Optional[Union[Tuple[int, int], int]] = None,
      padding: Union[
          ConvPadding, str, ConvPaddingInt
      ] = ConvPadding.VALID,
  ) -> JaxArray:
      """Applies max pooling using a square 2D filter.
  
      Args:
          x: input tensor of shape (N, C, H, W).
          size: size of pooling filter.
          strides: stride step, use size when stride is none (default).
          padding: padding of the input tensor, either Padding.SAME or Padding.VALID or numerical values.
  
      Returns:
          output tensor of shape (N, C, H, W).
      """
      size = to_tuple(size, 2)
      strides = to_tuple(strides, 2) if strides else size
      padding = to_padding(padding, 2)
      if isinstance(padding, tuple):
          padding = ((0, 0), (0, 0)) + padding
      return lax.reduce_window(
          x,
          -jn.inf,
          lax.max,
          (1, 1) + size,
          (1, 1) + strides,
          padding=padding,
      )
  
file_path: objax/objax/functional/core/pooling.py
imports: from typing import Union, Tuple, Optional
  import numpy
  from jax import numpy, lax
  from objax.constants import ConvPadding
  from objax.typing import JaxArray, ConvPaddingInt
  from objax.util import to_padding, to_tuple
label_desc: Use image node_id 4 for calling a global function with example usage: max_pool_2d(x, size, strides, padding) and returns: lax