code: def _run_state_jvp(
      primals: Sequence[Any],
      tangents: Sequence[Any],
      *,
      jaxpr: core.Jaxpr,
      which_linear: tuple[bool, ...],
  ):
      nonzero_tangents = [
          not isinstance(t, ad_util.Zero) for t in tangents
      ]
      discharged_jaxpr, body_consts = discharge_state(jaxpr, ())
      for _ in range(len(nonzero_tangents)):
          _, out_nonzero_tangents = ad.jvp_jaxpr(
              core.ClosedJaxpr(discharged_jaxpr, body_consts),
              nonzero_tangents,
              instantiate=nonzero_tangents,
          )
          if out_nonzero_tangents == nonzero_tangents:
              break
          nonzero_tangents = map(
              operator.or_, nonzero_tangents, out_nonzero_tangents
          )
      else:
          raise Exception("Invalid fixpoint")
      del discharged_jaxpr, body_consts, out_nonzero_tangents
      tangents = [
          ad.instantiate_zeros(t) if inst else t
          for t, inst in zip(tangents, nonzero_tangents)
      ]
      tangents = [t for t in tangents if type(t) is not ad_util.Zero]
      closed_jvp_jaxpr, _ = ad.jvp_jaxpr(
          pe.close_jaxpr(jaxpr), nonzero_tangents, []
      )
      jvp_jaxpr_, jvp_consts = (
          closed_jvp_jaxpr.jaxpr,
          closed_jvp_jaxpr.consts,
      )
      jvp_jaxpr = hoist_consts_to_refs(jvp_jaxpr_)
      jvp_which_linear = (
          *(False,) * len(jvp_consts),
          *which_linear,
          *(True,) * len(tangents),
      )
      out = run_state_p.bind(
          *jvp_consts,
          *primals,
          *tangents,
          jaxpr=jvp_jaxpr,
          which_linear=jvp_which_linear,
      )
      out_consts, out_primals, out_tangents = split_list(
          out, [len(jvp_consts), len(primals)]
      )
      del out_consts
      out_tangents_iter = iter(out_tangents)
      out_tangents = [
          next(out_tangents_iter) if nz else ad_util.Zero.from_value(p)
          for p, nz in zip(out_primals, nonzero_tangents)
      ]
      return out_primals, out_tangents
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import dataclasses
  from functools import partial
  import operator
  from typing import Any, Callable, Protocol
  import numpy
  from jax._src import api_util
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import tree_util
  from jax._src.interpreters import ad
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.lax import lax
  from jax._src.lax import slicing
  from jax._src.state import indexing
  from jax._src.state.types import AbstractRef, RefEffect
  from jax._src.state.primitives import get_p, swap_p, addupdate_p
  from jax._src.state.utils import hoist_consts_to_refs
  from jax._src.typing import Array
  from jax._src.util import safe_map, safe_zip, split_list, weakref_lru_cache, partition_list, merge_lists, split_dict
label_desc: Use image node_id 22 for calling a global function with example usage: _run_state_jvp(primals, tangents) and returns: out_primals,  out_tangents