code: def _jvp(f, x, b):
      """Computes jvp of `f` with respect to `x`."""
      broadcast_shape = ps.broadcast_shape(ps.shape(x), ps.shape(b))
      b = tf.broadcast_to(b, broadcast_shape)
      if JAX_MODE:
          import jax  # pylint:disable=g-import-not-at-top
  
          return jax.jvp(f.forward, (x,), (b,))[1]
      elif TF_MODE:
  
          @tf.function
          def jvp(b):
              with tf.autodiff.ForwardAccumulator(
                  primals=x, tangents=b
              ) as acc:
                  y = f.forward(x)
              result = acc.jvp(y)
              # TODO(srvasude): "result is None" might be caused by tf.Graph mode,
              # rather than missing the dependency of y w.r.t. x.
              if result is None:
                  result = tf.zeros_like(y)
              return result
  
          return jvp(b)
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 for calling a global function with example usage: _jvp(f, x, b) and returns: jax, jvp, result