newton_cg#

t3toolbox.backend.optimizers.newton_cg(problem, x0, max_newton=30, gtol_rel=1e-08, cg_maxiter=200, c_armijo=0.0001, g0norm_newton=None, g0norm_cg=None, cg_forcing_power=0.5, use_jit=False, callback=None)#
def newton_cg(
        problem:          Problem,    # the fixed-rank least-squares problem
        x0:               Tangent,    # initial cores (U, G)
        max_newton:       int   = 30,
        gtol_rel:         float = 1e-8,   # stop when ‖g‖ <= gtol_rel * ‖g0‖  (‖g0‖ = g0norm_newton or the initial ‖g‖)
        cg_maxiter:       int   = 200,
        c_armijo:         float = 1e-4,
        g0norm_newton:    typ.Optional[float] = None,  # ‖g0‖ for the Newton stop; default = initial ‖g‖. Also feeds CG unless g0norm_cg set.
        g0norm_cg:        typ.Optional[float] = None,  # ‖g0‖ for the CG forcing term; default = g0norm_newton (else initial ‖g‖)
        cg_forcing_power: float = 0.5,                 # η = min(0.5, (‖g‖/‖g0‖)**power); larger => tighter CG, fewer Newton steps
        use_jit:          bool  = False,  # jit the inner CG (lax.while_loop) when the inputs are jax; else eager
        callback:         typ.Optional[typ.Callable] = None,  # callback(NewtonInfo) each iteration (host-side; e.g. a display)
) -> typ.Tuple[Tangent, dict]:      # (x_cores, stats)

Inexact Riemannian Newton-CG with an Armijo line search – the manifold workhorse (the gauged H is positive-definite there). Each Newton step builds the local GN model once, solves H p = −g by CG to an inexact forcing-term tolerance (the inner loop – jit-able via _cg_solve()), then backtracks along retract(α p). The CG truncates on the gauge-singular corewise H; the outer line search keeps it robust regardless. use_jit jits only the inner CG (the outer loop, line search, and convergence test stay on the host), auto-converting numpy inputs to jax first (_prepare_jit_inputs()) – so a use_jit=True call on numpy data returns a jax-backed result (jax’s default float32 unless x64 is enabled) and raises if jax is not installed.

Overriding the reference gradient norm ‖g0‖ (g0norm_newton / g0norm_cg / cg_forcing_power). Both stopping tests are relative to a reference ‖g0‖: the Newton stop is ‖g‖ gtol_rel·‖g0‖ and the CG forcing term is η = min(0.5, (‖g‖/‖g0‖)**cg_forcing_power). By default ‖g0‖ is the initial gradient norm – but in a warm-start continuation loop that norm is misleadingly small (the guess is already near the solution), which over-tightens the Newton stop and slackens CG. Pass a reference reflecting the problem’s true gradient scale (e.g. the initial ‖g‖ from the first continuation stage) to restore the intended behavior. Resolution is a chained fallback: g0norm_newton sets the Newton reference (and CG inherits it unless g0norm_cg is also given); g0norm_cg alone sets only the CG reference. cg_forcing_power (default 0.5, the conventional Eisenstat-Walker value) tunes CG effort per Newton step: since ‖g‖/‖g0‖ < 1 near the solution, a larger power (0.75, 1.0) tightens CG -> more CG iterations but fewer Newton steps – worth it on the manifold when the retraction is expensive relative to a Hessian-apply. The min(0.5, …) cap on η is retained regardless.

callback, if given, is called with a NewtonInfo each iteration (including the final converged line) – the hook for a live diagnostic display; it runs host-side (it reads the concrete residual), so it composes with use_jit (only the inner CG jits) but not with a hypothetical fully-jitted outer loop. Ready-made displays: t3toolbox.backend.optimizer_display.make_newton_display(). stats always carries 'history' – one _newton_scalar_record() per iteration.

Parameters:
  • problem (Problem)

  • x0 (Tangent)

  • max_newton (int)

  • gtol_rel (float)

  • cg_maxiter (int)

  • c_armijo (float)

  • g0norm_newton (t3toolbox.backend.common.typ.Optional[float])

  • g0norm_cg (t3toolbox.backend.common.typ.Optional[float])

  • cg_forcing_power (float)

  • use_jit (bool)

  • callback (t3toolbox.backend.common.typ.Optional[t3toolbox.backend.common.typ.Callable])

Return type:

t3toolbox.backend.common.typ.Tuple[Tangent, dict]