tfp.experimental.substrates.jax.math.value_and_gradient

View source on GitHub

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