code: def _for_jvp(
      primals, tangents, *, jaxpr, nsteps, reverse, which_linear, unroll
  ):
      nonzero_tangents = [
          not isinstance(t, ad_util.Zero) for t in tangents
      ]
      # We need to find out which `Ref`s have nonzero tangents after running the
      # for loop. Ordinarily we do this with a fixed point on the body jaxpr but
      # a `for` body jaxpr is stateful and has no outputs. We therefore discharge
      # the state effect from the jaxpr and we will now have a "symmetric" jaxpr
      # where the inputs line up with the outputs. We use this discharged jaxpr
      # for the fixed point.
      closed_jaxpr = _cached_for_jaxpr(jaxpr)
      for _ in range(len(nonzero_tangents)):
          _, out_nonzero_tangents = ad.jvp_jaxpr(
              closed_jaxpr,
              [False] + 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")
      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_jaxpr = pe.close_jaxpr(jaxpr)
      jvp_jaxpr_, _ = ad.jvp_jaxpr(
          closed_jaxpr, [False] + nonzero_tangents, []
      )
      jvp_jaxpr, () = jvp_jaxpr_.jaxpr, jvp_jaxpr_.consts  # TODO consts
      jvp_which_linear = which_linear + (True,) * len(tangents)
      out_flat = for_p.bind(
          *primals,
          *tangents,
          jaxpr=jvp_jaxpr,
          nsteps=nsteps,
          reverse=reverse,
          which_linear=jvp_which_linear,
          unroll=unroll,
      )
      # `out_flat` includes constant inputs into the `for_loop` which are converted
      # into outputs as well. We don't care about these in AD so we throw them out.
      out_primals, out_tangents = split_list(out_flat, [len(primals)])
      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 functools
  import operator
  from typing import Any, Callable, Generic, TypeVar
  import jax.numpy
  from jax import lax
  from jax.api_util import flatten_fun_nokwargs
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax.tree_util import tree_flatten, tree_structure, tree_unflatten, treedef_tuple, tree_map, tree_leaves, PyTreeDef
  from jax._src import ad_util
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src.state.types import ReadEffect, AbstractRef, StateEffect
  from jax._src.state import discharge
  from jax._src.state import primitives
  from jax._src.state import utils
  from jax._src.state import types
  from jax._src.typing import Array
  from jax._src.util import partition_list, merge_lists, safe_map, safe_zip, split_list, split_dict, weakref_lru_cache
  from jax._src.lax.control_flow import loops
  from jax._src.lax.control_flow.common import _abstractify, _initial_style_jaxpr
label_desc: Use image node_id 11 for calling a global function with example usage: _for_jvp(primals, tangents) and returns: out_primals,  out_tangents