code: def _sample_many_kernel_fn(
      kernel_fn_sample_once,
      key: jax.Array,
      n_samples: set[int],
      get_generator: bool,
  ):
      def normalize(sample: PyTree, n: int) -> PyTree:
          return tree_map(lambda sample: sample / n, sample)
  
      def get_samples(
          x1: NTTree[jnp.ndarray],
          x2: Optional[NTTree[jnp.ndarray]],
          get: Get,
          **apply_fn_kwargs,
      ):
          _key = key
          ker_sampled = None
          for n in range(1, max(n_samples) + 1):
              _key, split = random.split(_key)
              one_sample = kernel_fn_sample_once(
                  x1, x2, split, get, **apply_fn_kwargs
              )
              if ker_sampled is None:
                  ker_sampled = one_sample
              else:
                  ker_sampled = tree_map(
                      operator.add, ker_sampled, one_sample
                  )
              yield n, ker_sampled
  
      if get_generator:
  
          @utils.get_namedtuple("MonteCarloKernel")
          def get_sampled_kernel(
              x1: jnp.ndarray,
              x2: jnp.ndarray,
              get: Optional[Get] = None,
              **apply_fn_kwargs,
          ) -> Generator[
              Union[jnp.ndarray, tuple[jnp.ndarray, ...]], None, None
          ]:
              for n, sample in get_samples(
                  x1, x2, get, **apply_fn_kwargs
              ):
                  if n in n_samples:
                      yield normalize(sample, n)
  
      else:
  
          @utils.get_namedtuple("MonteCarloKernel")
          def get_sampled_kernel(
              x1: jnp.ndarray,
              x2: jnp.ndarray,
              get: Optional[Get] = None,
              **apply_fn_kwargs,
          ) -> Union[jnp.ndarray, tuple[jnp.ndarray, ...]]:
              for n, sample in get_samples(
                  x1, x2, get, **apply_fn_kwargs
              ):
                  pass
              return normalize(sample, n)
  
      return get_sampled_kernel
  
imports: from functools import partial
  import operator
  from typing import Generator, Iterable, Optional, Union
  import jax
  from jax import random
  import jax.numpy
  from jax.tree_util import tree_map
  from .batching import batch
  from .empirical import _DEFAULT_NTK_FWD
  from .empirical import _DEFAULT_NTK_J_RULES
  from .empirical import _DEFAULT_NTK_S_RULES
  from .empirical import DEFAULT_NTK_IMPLEMENTATION
  from .empirical import empirical_kernel_fn
  from .empirical import NtkImplementation
  from .utils import utils
  from .utils.typing import ApplyFn
  from .utils.typing import Axes
  from .utils.typing import EmpiricalGetKernelFn
  from .utils.typing import Get
  from .utils.typing import InitFn
  from .utils.typing import MonteCarloKernelFn
  from .utils.typing import NTTree
  from .utils.typing import PyTree
  from .utils.typing import VMapAxes
label_desc: Use image node_id 2 for calling a global function with example usage: _sample_many_kernel_fn(kernel_fn_sample_once, key, n_samples, get_generator) and returns: get_sampled_kernel, tree_map, normalize