pypomp.Newton¶

class pypomp.Newton(clip_norm: float | None = None, scale: bool = False, ls: bool = False, c: float = 0.1, max_ls_itn: int = 10)[source]¶

Bases: Optimizer

Classic Second-Order Newton-Raphson optimizer.

init_state(theta: Array) → tuple[source]¶

Initialize the optimizer state buffers.

Parameters:: theta (jax.Array) – The parameters array. Can be 1D for standard parameters or ND for chunked unit parameters.
Returns:: The initial optimizer state buffers as a tuple of JAX arrays.
Return type:: tuple

step(grad: Array, state: tuple, step_num: int | Array, compute_hessian_fn: Callable[[], Array] | None = None, eta_i: Array | None = None) → tuple[Array, tuple][source]¶

Compute the parameter update direction and update the optimizer state.

Parameters:

grad (jax.Array) – The gradient of the objective function with respect to the parameters.
state (tuple) – The current optimizer state buffers.
step_num (int or jax.Array) – The current iteration or step index (0-indexed).
compute_hessian_fn (callable, optional) – A zero-argument callable that computes the model Hessian when invoked. Only called if the optimizer requires the Hessian (e.g. Newton methods).
eta_i (jax.Array, optional) – The current step size/learning rate at this iteration. Only used if the optimizer requires it (e.g. BFGS).

Returns:

direction (jax.Array) – The parameter update search direction.
new_state (tuple) – The updated optimizer state buffers.