code: def supports_masking(remask_kernel: bool):
      """Returns a decorator that turns layers into layers supporting masking.
  
      Specifically:
  
      1. `init_fn` is left unchanged.
  
      2. `apply_fn` is turned from a function that accepts a `mask=None` keyword
      argument (which indicates `inputs[mask]` must be masked), into a function
      that accepts a `mask_constant=None` keyword argument (which indicates
      `inputs[inputs == mask_constant]` must be masked).
  
      3. `kernel_fn` is modified to
  
        3.a. propagate the `kernel.mask1` and `kernel.mask2` through intermediary
        layers, and,
  
        3.b. if `remask_kernel == True`, zeroes-out covariances between entries of
        which at least one is masked.
  
      4. If the decorated layers has a `mask_fn`, it is used to propagate masks
      forward through the layer, in both `apply_fn` and `kernel_fn`. If not, it is
      assumed the mask remains unchanged.
  
      Must be applied before the `layer` decorator.
  
      See Also:
        Example of masking application in `examples/imdb.py`.
  
      Args:
        remask_kernel:
          `True` to zero-out kernel covariance entries between masked inputs after
          applying `kernel_fn`. Some layers don't need this and setting
          `remask_kernel=False` can save compute.
  
      Returns:
        A decorator that turns functions returning
        `(init_fn, apply_fn, kernel_fn[, mask_fn])`
        into functions returning
        `(init_fn, apply_fn_with_masking, kernel_fn_with_masking)`.
      """
  
      def supports_masking(layer):
          @utils.wraps(layer)
          def layer_with_masking(*args, **kwargs) -> InternalLayer:
              layer_fns = layer(*args, **kwargs)
              init_fn, apply_fn, kernel_fn = layer_fns[:3]
  
              if len(layer_fns) == 3:
                  # No mask propagation function supplied - use identity.
                  _mask_fn = lambda mask, input_shape: mask
              elif len(layer_fns) == 4:
                  # Custom mask propagation function supplied.
                  _mask_fn = layer_fns[3]
              else:
                  raise ValueError(
                      f"Expected 3 (`init_fn`, `apply_fn`, `kernel_fn`) or 4"
                      f" (..., `mask_fn`) layer functions, "
                      f"got {len(layer_fns)}."
                  )
  
              @utils.wraps(_mask_fn)
              def mask_fn(mask, input_shape):
                  if mask is None:
                      return None
                  return _mask_fn(mask, input_shape)
  
              def apply_fn_with_masking(
                  params, inputs, *, mask_constant=None, **kwargs
              ):
                  masked_inputs = tree_map(
                      lambda x: _get_masked_array(x, mask_constant),
                      inputs,
                      is_leaf=lambda x: isinstance(
                          x, (jnp.ndarray, MaskedArray)
                      ),
                  )
  
                  is_leaf = lambda x: isinstance(x, MaskedArray)
                  inputs = tree_map(
                      lambda x: x.masked_value,
                      masked_inputs,
                      is_leaf=is_leaf,
                  )
                  mask = tree_map(
                      lambda x: x.mask, masked_inputs, is_leaf=is_leaf
                  )
  
                  outputs = apply_fn(
                      params, inputs, mask=mask, **kwargs
                  )
                  outputs_mask = mask_fn(
                      mask,
                      inputs.shape
                      if isinstance(inputs, jnp.ndarray)
                      else [i.shape for i in inputs],
                  )
                  if outputs_mask is None:
                      return outputs
                  return MaskedArray(
                      outputs, outputs_mask
                  )  # pytype:disable=wrong-arg-count
  
              def kernel_fn_with_masking(
                  k: NTTree[Kernel], **user_reqs
              ):
                  is_leaf = lambda k: isinstance(k, Kernel)
                  mask1 = tree_map(
                      lambda k: k.mask1, k, is_leaf=is_leaf
                  )
                  shape1 = tree_map(
                      lambda k: k.shape1, k, is_leaf=is_leaf
                  )
                  mask2 = tree_map(
                      lambda k: k.mask2, k, is_leaf=is_leaf
                  )
                  shape2 = tree_map(
                      lambda k: k.shape2, k, is_leaf=is_leaf
                  )
  
                  mask1, mask2 = mask_fn(mask1, shape1), mask_fn(
                      mask2, shape2
                  )
  
                  k = kernel_fn(k, **user_reqs)  # type: Kernel
  
                  if remask_kernel:
                      remask_fn = lambda k, m1, m2: k.mask(m1, m2)
                  else:
                      remask_fn = lambda k, m1, m2: k.replace(
                          mask1=m1, mask2=m2
                      )
  
                  k = tree_map(
                      remask_fn, k, mask1, mask2, is_leaf=is_leaf
                  )
                  return k
  
              if _has_req(kernel_fn):
                  _set_req(kernel_fn_with_masking, get_req(kernel_fn))
  
              return (
                  init_fn,
                  apply_fn_with_masking,
                  kernel_fn_with_masking,
              )
  
          return layer_with_masking
  
      return supports_masking
  
imports: import dataclasses
  import enum
  from typing import Callable, Optional, Sequence, Union
  import warnings
  import frozendict
  import jax
  from jax import eval_shape
  from jax import lax
  from jax import numpy
  from jax.core import ShapedArray
  from jax.tree_util import tree_all
  from jax.tree_util import tree_map
  import numpy
  from ..utils import dataclasses
  from ..utils import utils
  from ..utils.kernel import Kernel
  from ..utils.typing import AnalyticKernelFn
  from ..utils.typing import ApplyFn
  from ..utils.typing import Axes
  from ..utils.typing import Get
  from ..utils.typing import InitFn
  from ..utils.typing import InternalLayer
  from ..utils.typing import Layer
  from ..utils.typing import LayerKernelFn
  from ..utils.typing import NTTree
  from ..utils.typing import PyTree
label_desc: Use image node_id 3 for calling a global function with example usage: supports_masking(remask_kernel) and returns: supports_masking, layer_with_masking, init_fn,  apply_fn_with_masking,  kernel_fn_with_masking, _mask_fn, MaskedArray, k, None, outputs