code: def _pallas_call_jvp_rule(
      primals,
      tangents,
      *,
      jaxpr,
      name,
      which_linear,
      input_output_aliases: tuple[tuple[int, int], ...],
      in_shapes,
      out_shapes,
      grid_mapping,
      debug,
      interpret,
      **compiler_params: Any,
  ):
      if grid_mapping.num_index_operands:
          raise NotImplementedError
      if input_output_aliases:
          raise NotImplementedError("JVP with aliasing not supported.")
      nonzero_tangents = [
          not isinstance(t, ad_util.Zero) for t in tangents
      ]
      tangents = [t for t in tangents if type(t) is not ad_util.Zero]
      nonzero_tangents_with_outputs = nonzero_tangents + [True] * len(
          out_shapes
      )
      closed_jaxpr = jax_core.ClosedJaxpr(jaxpr, ())
      jvp_jaxpr_, _ = ad.jvp_jaxpr(
          closed_jaxpr, nonzero_tangents_with_outputs, []
      )
      jvp_jaxpr, () = jvp_jaxpr_.jaxpr, jvp_jaxpr_.consts  # TODO consts
      # `pallas_call` takes in inputs and returns outputs but its jaxpr *does not*.
      # `pallas_call` takes in a stateful jaxpr, meaning the jaxpr accepts input
      # `Ref`s that are read from followed by output `Ref`s that are written to.
      # This means that when we do `jvp_jaxpr` on the `jaxpr`, we get out a new
      # jaxpr that has tangents following primals. In order for this jaxpr to be
      # compatible w/ `pallas_call` (inputs then outputs), we need to shuffle around
      # the jaxpr's invars.
      (
          primal_refs,
          primal_out_refs,
          tangent_refs,
          tangent_out_refs,
      ) = split_list(
          jvp_jaxpr.invars,
          [len(primals), len(out_shapes), len(tangents)],
      )
      invars = (
          *primal_refs,
          *tangent_refs,
          *primal_out_refs,
          *tangent_out_refs,
      )
      # TODO(sharadmv): Fix state effect tracking after invar switch.
      jvp_jaxpr = jvp_jaxpr.replace(invars=invars)
      if debug:
          print(jvp_jaxpr)
      in_bms, out_bms = split_list(
          grid_mapping.block_mappings, [len(primals)]
      )
      jvp_bms = (*in_bms, *in_bms, *out_bms, *out_bms)
      out_flat = pallas_call_p.bind(
          *primals,
          *tangents,
          jaxpr=jvp_jaxpr,
          name=f"{name}_jvp",
          in_shapes=(*in_shapes, *in_shapes),
          out_shapes=(*out_shapes, *out_shapes),
          grid_mapping=grid_mapping.replace(block_mappings=jvp_bms),
          which_linear=which_linear + (True,) * len(tangents),
          interpret=interpret,
          debug=debug,
          input_output_aliases=(),
          **compiler_params,
      )
      out_primals, out_tangents = split_list(
          out_flat, [len(out_flat) // 2]
      )
      return out_primals, out_tangents
  
imports: from __future__ import annotations
  from functools import partial
  import itertools
  from typing import Any, Callable
  from collections.abc import Sequence
  import jax
  from jax import api_util
  from jax import tree_util
  from jax import lax
  from jax._src import state
  from jax.interpreters import ad
  from jax.interpreters import batching
  from jax.interpreters import partial_eval
  from jax.interpreters import mlir
  from jax.interpreters import xla
  from jax._src import ad_util
  from jax._src import core
  from jax._src.state import primitives
  from jax._src import linear_util
  from jax._src.state import discharge
  from jax._src.util import split_list, safe_map, safe_zip, weakref_lru_cache, tuple_insert, partition_list, merge_lists
  import jax.numpy
  import numpy
  from jax._src.pallas import core
label_desc: Use image node_id 6 for calling a global function with example usage: _pallas_call_jvp_rule(primals, tangents) and returns: out_primals,  out_tangents