code: def get_if_single_primitive(
      fun: Callable[..., Any], *args: Any
  ) -> jax.core.Primitive | None:
      """
      If fun(*args) lowers to a single primitive with inputs and outputs matching
      function inputs and outputs, return that primitive. Otherwise return None.
      """
      try:
          jaxpr = jax.make_jaxpr(fun)(*args)
      except:
          return None
      while len(jaxpr.eqns) == 1:
          eqn = jaxpr.eqns[0]
          if (eqn.invars, eqn.outvars) != (
              jaxpr.jaxpr.invars,
              jaxpr.jaxpr.outvars,
          ):
              return None
          elif eqn.primitive == jax._src.pjit.pjit_p and all(
              jax._src.pjit.is_unspecified(sharding)
              for sharding in (
                  *eqn.params["in_shardings"],
                  *eqn.params["out_shardings"],
              )
          ):
              jaxpr = jaxpr.eqns[0].params["jaxpr"]
          else:
              return jaxpr.eqns[0].primitive
      return None
  
imports: from __future__ import annotations
  from functools import partial
  import math
  import operator
  from typing import Any, Callable
  import jax
  from jax._src.typing import Array, ArrayLike, DTypeLike
  from jax._src.lax import lax
  from jax._src.numpy import reductions
  from jax._src.numpy.lax_numpy import _eliminate_deprecated_list_indexing, append, take
  from jax._src.numpy.reductions import _moveaxis
  from jax._src.numpy.util import _wraps, check_arraylike, _broadcast_to, _where
  from jax._src.numpy.vectorize import vectorize
  from jax._src.util import canonicalize_axis, set_module
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: get_if_single_primitive(fun) and returns: None, None, None, jaxpr