code: def unmask_fn(fn: ApplyFn) -> ApplyFn:
      """Make a function returning a `MaskedArray` return a `jnp.ndarray`.
  
      Useful if you pass `masked_constant` to your `apply_fn` in order to have
      variable-length inputs. In this case `apply_fn` returns a `MaskedArray`
      that stores the information about which entries are masked (for convenient
      chaining with further functions operating on masked inputs). This decorator
      replaces the output `MaskedArray` with an `jnp.ndarray` where masked
      entries are zeroed-out, which is convenient to pass to functions operating on
      arrays, such as :obj:`~neural_tangents.monte_carlo_kernel_fn` or
      :obj:`~neural_tangents.empirical_kernel_fn`.
  
      .. warning::
        In some cases you may want to define your own custom unmasking behavior,
        e.g. one that normalizes the values based on the number of non-zero entries.
  
      See Also:
        :class:`MaskedArray`, and an example masking application in
        `examples/imdb.py`.
  
      Args:
        fn: function returning a :class:`MaskedArray`.
  
      Returns:
        Function of same signature as `fn`, where the output :class:`MaskedArray` is
        replaced with the :class:`jax.numpy.ndarray` with masked entries zeroed-out.
      """
  
      def unmask(x: Union[MaskedArray, jnp.ndarray]) -> jnp.ndarray:
          if isinstance(x, MaskedArray):
              x = utils.mask(x.masked_value, x.mask)
          return x  # pytype: disable=bad-return-type  # jax-ndarray
  
      def is_leaf(x) -> bool:
          return isinstance(x, (jnp.ndarray, MaskedArray))
  
      @utils.wraps(fn)
      def fn_no_mask(*args, **kwargs):
          out = fn(*args, **kwargs)
          out = tree_map(unmask, out, is_leaf=is_leaf)
          return out
  
      return fn_no_mask
  
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 4 for calling a global function with example usage: unmask_fn(fn) and returns: fn_no_mask, x, isinstance, out