code: def max_pool2d(
      x: JaxArray,
      kernel: Union[int, Tuple[int, ...]],
      strides: Union[int, Tuple[int, ...]],
      padding: Union[str, int, Tuple[int], List[Tuple[int, int]]],
      /,
      *,
      data_format: str = "NHWC",
      dilation: Union[int, Tuple[int, ...]] = 1,
      ceil_mode: bool = False,
      out: Optional[JaxArray] = None,
  ) -> JaxArray:
      dims = 2
      odtype = x.dtype
      kernel, strides, padding, dilation = _validate_max_pool_params(
          kernel,
          strides,
          padding,
          dilation,
          ceil_mode,
          dims,
          data_format,
      )
  
      if data_format == "NCHW":
          x = jnp.transpose(x, (0, 2, 3, 1))
          kernel = (
              [kernel[i] for i in [0, 2, 3, 1]]
              if len(kernel) == (dims + 2)
              else kernel
          )
          strides = (
              [strides[i] for i in [0, 2, 3, 1]]
              if len(strides) == (dims + 2)
              else strides
          )
          padding = (
              [padding[i] for i in [0, 2, 3, 1]]
              if isinstance(padding, list)
              and len(padding) == (dims + 2)
              else padding
          )
  
      res = general_pool(
          x,
          -jnp.inf,
          jlax.max,
          kernel,
          strides,
          padding,
          dims,
          dilation,
          ceil_mode,
      )
  
      if data_format == "NCHW":
          res = jnp.transpose(res, (0, 3, 1, 2))
  
      return res.astype(odtype)
  
imports: from typing import Optional, Union, Tuple, List, Literal, Sequence, Callable
  import jax
  import jax.lax
  import jax.numpy
  import math
  import ivy
  from ivy import output_to_native_arrays
  from ivy.functional.backends.jax import JaxArray
  from ivy.functional.backends.jax.random import RNG
  from ivy.functional.ivy.experimental.general import _correct_ivy_callable
  from ivy.functional.ivy.layers import _handle_padding, _validate_max_pool_params, _depth_max_pooling_helper
  from ivy.functional.ivy.experimental.layers import _padding_ceil_mode, _get_size
  from ivy.func_wrapper import with_supported_dtypes
  from ivy.func_wrapper import with_unsupported_dtypes
  from .None import backend_version
  from ivy.functional.backends.jax.experimental.manipulation import _to_nested_tuple
label_desc: Use image node_id 5 for calling a global function with example usage: max_pool2d() and returns: res