code: def _log_bessel_kve_jvp(primals, tangents):
      """Computes JVP for bessel_kve (supports JAX custom derivative)."""
      v, z = primals
      _, dz = tangents
  
      dtype = dtype_util.common_dtype([v, z], tf.float32)
      numpy_dtype = dtype_util.as_numpy_dtype(dtype)
  
      # 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)
  
      log_kve = _log_bessel_kve_custom_gradient(v, z)
      pz = (
          tfp_math.log_add_exp(
              _log_bessel_kve_custom_gradient(v - 1.0, z),
              _log_bessel_kve_custom_gradient(v + 1.0, z),
          )
          - numpy_dtype(np.log(2.0))
          - log_kve
      )
      pz = -tf.math.expm1(pz)
  
      # `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 log_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 38 for calling a global function with example usage: _log_bessel_kve_jvp(primals, tangents) and returns: log_kve,  unknown