code: def _bessel_kve_jvp(primals, tangents):
      """Computes JVP for bessel_kve (supports JAX custom derivative)."""
      (
          v,
          z,
      ) = primals
      _, dz = tangents
  
      # TODO(https://github.com/google/jax/issues/3768): eliminate broadcast_to?
      bc_shp = ps.broadcast_shape(ps.shape(v), ps.shape(dz))
      dz = tf.broadcast_to(dz, bc_shp)
  
      kve = _bessel_kve_custom_gradient(v, z)
      pz = (z - v) / z * kve - _bessel_kve_custom_gradient(v - 1.0, z)
  
      # `bessel_kve` does not have gradients with respect to `v`, and thus
      # this `JVP` rule matches TF.
      # Ideally, it would be nice to throw an exception when taking gradients of
      # in JAX mode, but this is not possible at the moment with `custom_jvp`.
      # See https://github.com/google/jax/issues/5913 for details.
      # TODO(https://github.com/google/jax/issues/5913): Define vjp for v.
  
      return kve, pz * dz
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math import generic
label_desc: Use image node_id 26 for calling a global function with example usage: _bessel_kve_jvp(primals, tangents) and returns: kve,  unknown