code: def try_import_tfp_jax_op(
      op: ScalarOp, jax_op_name: Optional[str] = None
  ) -> Callable:
      try:
          import tensorflow_probability.substrates.jax.math as tfp_jax_math
      except ModuleNotFoundError:
          raise NotImplementedError(
              f"No JAX implementation for Op {op.name}. "
              "Implementation is available if TensorFlow Probability is installed"
          )
  
      if jax_op_name is None:
          jax_op_name = op.name
      return typing.cast(Callable, getattr(tfp_jax_math, jax_op_name))
  
imports: import functools
  import typing
  from typing import Callable, Optional
  import jax
  import jax.numpy
  from pytensor.link.jax.dispatch.basic import jax_funcify
  from pytensor.scalar import Softplus
  from pytensor.scalar.basic import Add, Cast, Clip, Composite, Identity, IntDiv, Mod, Mul, ScalarOp, Second, Sub
  from pytensor.scalar.math import BetaIncInv, Erf, Erfc, Erfcinv, Erfcx, Erfinv, GammaIncCInv, GammaIncInv, Iv, Ive, Log1mexp, Psi, TriGamma
label_desc: Use image node_id 1 for calling a global function with example usage: try_import_tfp_jax_op(op, jax_op_name) and returns: typing