code: def _jvp_transform(static_fn):
      def _jvp_transform_impl(args):
          args, t_args = args
          tang_is_none = jtu.tree_map(
              _is_none, t_args, is_leaf=_is_none
          )
          no_tang, has_tang = partition(
              args, tang_is_none, is_leaf=_is_none
          )
  
          def _to_jvp(_has_tang):
              _args = combine(_has_tang, no_tang)
              return static_fn(_args)
  
          _, t_out = filter_jvp(_to_jvp, (has_tang,), (t_args,))
          return t_out
  
      return _jvp_transform_impl
  
imports: import functools
  from collections.abc import Callable
  from typing import Any, Optional, Union
  import jax
  import jax._src.pretty_printer
  import jax._src.source_info_util
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import Array, Int, PyTree
  from .._ad import filter_jvp
  from .._caches import internal_lru_caches
  from .._compile_utils import hashable_combine, hashable_partition
  from .._eval_shape import filter_eval_shape
  from .._filters import combine, is_array, partition
  from .._module import field, Module, module_update_wrapper
  from .._vmap_pmap import filter_vmap
  from .None import _primitive
  from ._primitive import filter_primitive_batching, filter_primitive_bind, filter_primitive_def, filter_primitive_jvp, filter_primitive_transpose, materialise_zeros
label_desc: Use image node_id 13 for calling a global function with example usage: _jvp_transform(static_fn) and returns: _jvp_transform_impl, t_out, static_fn