code: def _get_vjp_fun(
      primal_fun: Callable,
      *,
      in_tree: tree_util.PyTreeDef,
      in_avals: Sequence[core.AbstractValue],
      out_avals: Sequence[core.AbstractValue],
      in_shardings: tuple[Sharding, ...],
      out_shardings: tuple[Sharding, ...],
      nr_devices: int,
      apply_jit: bool,
  ) -> tuple[Callable, Sequence[core.AbstractValue]]:
      # Since jax.vjp does not handle kwargs, it is easier to do all the work
      # here with flattened functions.
      def fun_vjp_jax(*args_and_out_cts_flat_jax):
          # Takes a flat list of primals and output cotangents
          def flattened_primal_fun_jax(*args_flat):
              args, kwargs = in_tree.unflatten(args_flat)
              res = primal_fun(*args, **kwargs)
              res_flat, _ = tree_util.tree_flatten(res)
              return res_flat
  
          args_flat_jax, out_cts_flat_jax = util.split_list(
              args_and_out_cts_flat_jax, [len(in_avals)]
          )
          _, pullback_jax = jax.vjp(
              flattened_primal_fun_jax, *args_flat_jax
          )
          return pullback_jax(out_cts_flat_jax)
  
      vjp_in_avals = list(
          itertools.chain(
              in_avals, map(lambda a: a.at_least_vspace(), out_avals)
          )
      )
  
      if apply_jit:
          # Prepare a device assignment. For exporting purposes, all it matters
          # is the number of devices.
          device_assignment = jax.devices(jax.default_backend())[
              :nr_devices
          ]
          assert len(device_assignment) == nr_devices
          vjp_in_shardings, vjp_out_shardings = canonical_shardings(
              device_assignment,
              tuple(itertools.chain(in_shardings, out_shardings)),
              in_shardings,
          )
          return (
              pjit.pjit(
                  fun_vjp_jax,
                  in_shardings=vjp_in_shardings,
                  out_shardings=vjp_out_shardings,
              ),
              vjp_in_avals,
          )
      else:
          return fun_vjp_jax, vjp_in_avals
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import copy
  import dataclasses
  import functools
  import itertools
  import re
  from typing import Any, Callable, Union
  import warnings
  from absl import logging
  import numpy
  import jax
  from jax import sharding
  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.interpreters import mlir
  from jax._src.interpreters import pxla
  from jax._src.lib import xla_client
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import hlo
  from jax._src.lib.mlir.dialects import func
  from jax._src import pjit
  from jax._src import sharding_impls
  from jax._src import source_info_util
  from jax._src import tree_util
  from jax._src import util
  from jax._src import xla_bridge
  from jax.experimental.export import _shape_poly
label_desc: Use image node_id 13 for calling a global function with example usage: _get_vjp_fun(primal_fun) and returns: pullback_jax, pjit,  vjp_in_avals, fun_vjp_jax,  vjp_in_avals, res_flat