code: def _sample_once_kernel_fn(
      kernel_fn: EmpiricalGetKernelFn,
      init_fn: InitFn,
      batch_size: int = 0,
      device_count: int = -1,
      store_on_device: bool = True,
  ):
      @partial(
          batch,
          batch_size=batch_size,
          device_count=device_count,
          store_on_device=store_on_device,
      )
      def kernel_fn_sample_once(
          x1: NTTree[jnp.ndarray],
          x2: Optional[NTTree[jnp.ndarray]],
          key: jax.Array,
          get: Get,
          **apply_fn_kwargs,
      ):
          init_key, dropout_key = random.split(key, 2)
          shape = tree_map(lambda x: x.shape, x1)
          _, params = init_fn(init_key, shape)
          return kernel_fn(
              x1, x2, get, params, rng=dropout_key, **apply_fn_kwargs
          )
  
      return kernel_fn_sample_once
  
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 1 for calling a global function with example usage: _sample_once_kernel_fn(kernel_fn, init_fn, batch_size, device_count, store_on_device) and returns: kernel_fn_sample_once, kernel_fn