newton_cg
=========

.. py:function:: 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)

   .. code-block:: python

      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 :py:func:`_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
   (:py:func:`_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 :py:class:`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: :py:func:`t3toolbox.backend.optimizer_display.make_newton_display`.
   ``stats`` always carries ``'history'`` -- one :py:func:`_newton_scalar_record` per iteration.