code: def _cond_jvp(primals, tangents, branches, linear):
      nonzeros = [type(t) is not ad_util.Zero for t in tangents]
  
      index_nz, *ops_nz = nonzeros
      assert index_nz is False
  
      branches_out_nz = [
          ad.jvp_jaxpr(jaxpr, ops_nz, instantiate=False)[1]
          for jaxpr in branches
      ]
      out_nz = [any(nz) for nz in zip(*branches_out_nz)]
  
      branches_jvp = tuple(
          ad.jvp_jaxpr(jaxpr, ops_nz, instantiate=out_nz)[0]
          for jaxpr in branches
      )
  
      index, *ops = primals
      _, *ops_dot = tangents
      ops_dot = _prune_zeros(ops_dot)
  
      ops_lin = tuple(linear)
      linear_jvp = ops_lin + (True,) * len(ops_dot)
      out = cond_p.bind(
          index,
          *ops,
          *ops_dot,
          branches=branches_jvp,
          linear=linear_jvp,
      )
      out_primals, out_tangents = split_list(out, [len(out_nz)])
      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, out_nz)
      ]
      return out_primals, out_tangents
  
imports: from __future__ import annotations
  import collections
  from collections.abc import Sequence
  import functools
  from functools import partial
  import inspect
  import itertools
  import operator
  from typing import Any, Callable, TypeVar
  from jax.tree_util import tree_flatten, tree_unflatten
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import effects
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import util
  from jax._src.state.discharge import register_discharge_rule, discharge_state
  from jax._src.state.types import AbstractRef, RefEffect
  from jax._src.core import ConcreteArray, raise_to_shaped, replace_jaxpr_effects
  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._src.interpreters import xla
  from jax._src.lax import lax
  from jax._src.traceback_util import api_boundary
  from jax._src.util import safe_map, split_list, partition_list
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import hlo
  import numpy
  from jax._src.lax.control_flow.common import _abstractify, _avals_short, _check_tree_and_avals, _initial_style_jaxprs_with_common_consts, _make_closed_jaxpr, _prune_zeros, _typecheck_param
label_desc: Use image node_id 10 for calling a global function with example usage: _cond_jvp(primals, tangents, branches, linear) and returns: out_primals,  out_tangents