hvp_rad#
- scarlet2.infer.hvp_rad(hvp, shape, key=None, max_iter=100)[source]#
Approximate the diagonal of the Hessian
Uses the Hutchinson’s method to approximate the diagonal of the Hessian by computing the Hessian-vector products with random Rademacher vectors.
- Parameters:
hvp (
callable) – A function that computes the Hessian-vector product.shape (
tuple) – The shape of the input to the Hessian-vector product function.key (
jax.random.PRNGKey, optional) – A random key for generating Rademacher vectors. If None, a new key will be generated.max_iter (
int, optional) – The maximum number of iterations for the approximation. Default is 100.
- Returns:
An array containing the approximate diagonal of the Hessian.
- Return type:
jax.numpy.ndarray