ad_utils

jaxns.experimental.solvers.ad_utils

Module Contents

tree_dot(x, y)

tree_norm(x)

tree_mul(x, y)

tree_div(x, y)

tree_vdot(x, y)

tree_vdot_real_part(x, y)

tree_add(x, y)

tree_scalar_mul(alpha, x)

tree_neg(x)

tree_sub(x, y)

hvp_linearized(f, params)

hvp_forward_over_reverse(f, params)

hvp_reverse_over_reverse(f, params)

hvp_reverse_over_forward(f, params)

grad_and_hvp(f, params, v)

Compute the gradient and Hessian-vector product of a function.

Parameters:

• f – the function to differentiate, should be scalar output

• params – the parameters to differentiate with respect to

• v – the vector to multiply the Hessian with

Returns:

the gradient and Hessian-vector product

build_hvp(f, params, linearise=True)

Build a function that computes the Hessian-vector product of a function.

Parameters:

• f – scalar function to differentiate

• params – the parameters to differentiate with respect to

• linearise (bool) – whether to linearize the gradient function at params, can be better for reapplying the HVP multiple times.

Returns:

a function that computes the Hessian-vector product