code: def _load_jvp(primals, tangents, args_tree, **params):
      ref_primal, indexers, mask, other_primal = args_tree.unflatten(
          primals
      )
      ref_tangent, _, _, other_tangent = args_tree.unflatten(tangents)
      if other_tangent is not None:
          other_tangent = ad_util.instantiate(other_tangent)
      return (
          load_p.bind(
              *tree_util.flatten(
                  ref_primal, indexers, mask, other_primal
              ),
              args_tree=args_tree,
              **params,
          ),
          load_p.bind(
              *tree_util.flatten(
                  ref_tangent, indexers, mask, other_tangent
              ),
              args_tree=args_tree,
              **params,
          ),
      )
  
imports: from __future__ import annotations
  import enum
  import functools
  from typing import Any
  import jax
  from jax import lax
  from jax import tree_util
  from jax._src import ad_util
  from jax._src import core
  from jax._src import pretty_printer
  from jax._src import state
  from jax._src.util import safe_map, safe_zip
  from jax._src.state import discharge
  from jax._src.state import indexing
  from jax._src.state import primitives
  from jax.interpreters import ad
  from jax.interpreters import mlir
  from jax.interpreters import xla
  import jax.numpy
  from jax._src.pallas import core
label_desc: Use image node_id 17 for calling a global function with example usage: _load_jvp(primals, tangents, args_tree) and returns: load_p,  load_p