code: def dropout2d(
      x: JaxArray,
      prob: float,
      /,
      *,
      training: bool = True,
      data_format: str = "NHWC",
      out: Optional[JaxArray] = None,
  ) -> JaxArray:
      if training:
          x_shape = x.shape
          is_batched = len(x.shape) == 4
          if data_format == "NCHW":
              perm = (0, 2, 3, 1) if is_batched else (1, 2, 0)
              x = jnp.transpose(x, perm)
              x_shape = x.shape
          _, rng_input = jax.random.split(RNG.key)
          mask = jax.random.bernoulli(rng_input, 1 - prob, x_shape)
          res = jnp.where(mask, x / (1 - prob), 0)
          if data_format == "NCHW":
              perm = (0, 3, 1, 2) if is_batched else (2, 0, 1)
              res = jnp.transpose(res, perm)
      else:
          res = x
      return res
  
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 14 for calling a global function with example usage: dropout2d() and returns: res