Missed TensorFlow Dev Summit? Check out the video playlist. Watch recordings

tfp.experimental.substrates.jax.math.value_and_gradient

View source on GitHub

Computes f(*xs) and its gradients wrt to *xs.

tfp.experimental.substrates.jax.math.value_and_gradient(
    f, xs, output_gradients=None, use_gradient_tape=False, name=None
)