newton_cg#

t3toolbox.optimizers.newton_cg(geometry, kind, sample, data, x0, order=None, weight=None, verbose=False, val_sample=None, val_data=None, callback=None, use_jit=False, regularizer=None, chunk_size='auto', **kwargs)#
def newton_cg(
        geometry,                       # ragged/uniform MANIFOLD (intended) / COREWISE (must match x0)
        kind:     str,                  # 'apply' / 'entries' / 'probe' (+ '_derivatives')
        sample:   typ.Any,              # ww / index, or (ww, pp) / (index, pp) for derivatives
        data:     typ.Any,             # observed values to fit
        x0:       Point,                # initial point (zero is fine on the manifold)
        order:    typ.Optional[int] = None,
        weight:   typ.Optional[typ.Any] = None,   # residual weight ω: per-mode (probe) / ω[mode,order] (derivatives)
        verbose:  bool = False,                   # print per-iteration diagnostics (a relative-error table + more)
        val_sample: typ.Any = None,               # optional validation sample (same layout as `sample`) -> a train|val table
        val_data:   typ.Any = None,               # optional validation data (both given adds the val column)
        callback:   typ.Optional[typ.Callable] = None,  # custom callback(NewtonInfo) each iter (overrides `verbose`)
        use_jit:    bool = False,               # jit the inner CG: auto-converts x0/sample/data to jax -> a jax-backed float32 result; raises if jax absent
        regularizer: typ.Any = None,            # optional regularizer, e.g. optimizers.IdentityRegularizer(λ) (ragged only)
        chunk_size: typ.Any = 'auto',   # probe_derivatives 𝒥ᵀ memory chunk; 'auto' -> estimate_chunk_size (docs/chunking.md)
        **kwargs,                       # forwarded to backend.optimizers.newton_cg (max_newton, gtol_rel, g0norm_newton, ...)
) -> typ.Tuple[Point, dict]:

Inexact Riemannian Newton-CG with an Armijo line search – the manifold workhorse. Ragged or uniform x0 (see gradient_descent()). See t3toolbox.backend.optimizers.newton_cg().

In a warm-start continuation loop the initial gradient norm ‖g0‖ is misleadingly small, which over-tightens the Newton stop and slackens CG; pass g0norm_newton / g0norm_cg (forwarded via **kwargs) to override the reference norm the two relative stopping tests use, and cg_forcing_power to trade more CG iterations per Newton step for fewer Newton steps (all detailed on the backend).

verbose=True prints a per-iteration diagnostic block (objective / gradient, CG stats, line search, and the per-(mode, order) relative-error table); pass val_sample / val_data to add a validation column. With a regularizer attached the objective is shown split as obj = misfit + reg (the ½‖ω⊙r‖² data misfit vs ρ(x)); both parts are on every record as misfit / regularization (the latter None when unregularized), in stats['history'] (always) and stats['diagnostics'] (when verbose). This is a thin convenience over the backend display – a raw-.data user builds the same callback with t3toolbox.backend.optimizer_display.make_newton_display() and passes it as callback= to t3toolbox.backend.optimizers.newton_cg(). A custom callback overrides verbose. Works on both the ragged and uniform layers (the uniform block_sumsq reduces the packed residual directly; validation data is packed automatically).

For a large-|W| uniform probe_derivatives fit, chunk_size='auto' (the default) sizes the 𝒥ᵀ gradient assembly’s memory automatically (via estimate_chunk_size()); pass an int / None to override. Other kinds and the ragged layer ignore it. See Memory chunking of the derivative transpose (chunk_size).

Examples

Recover a low-rank tensor from noiseless apply measurements on the uniform layer – a uniform x0 (with UNIFORM_MANIFOLD) returns a UniformTuckerTensorTrain (a zero start is fine on the manifold; unit-norm probe rows keep the least-squares well-conditioned):

>>> import numpy as np
>>> import t3toolbox.tucker_tensor_train as t3
>>> import t3toolbox.uniform_tucker_tensor_train as ut3
>>> import t3toolbox.uniform_manifold as ut3m
>>> import t3toolbox.optimizers as optimizers
>>> np.random.seed(0)
>>> A = t3.TuckerTensorTrain.randn((6, 7, 8), (2, 2, 2), (1, 2, 2, 1))
>>> ww = [np.random.randn(120, N) for N in (6, 7, 8)]
>>> ww = [w / np.linalg.norm(w, axis=1, keepdims=True) for w in ww]   # unit-norm rows
>>> b = A.apply(ww)
>>> x0 = ut3.UniformTuckerTensorTrain.from_t3(t3.TuckerTensorTrain.zeros((6, 7, 8), (2, 2, 2), (1, 2, 2, 1)))
>>> x_opt, stats = optimizers.newton_cg(ut3m.UNIFORM_MANIFOLD, 'apply', ww, b, x0, max_newton=30)
>>> type(x_opt).__name__                              # returned in the same (uniform) representation
'UniformTuckerTensorTrain'
>>> bool(np.linalg.norm(x_opt.to_dense() - A.to_dense()) / np.linalg.norm(A.to_dense()) < 1e-6)
True

For a ragged fit pass t3m.MANIFOLD and a TuckerTensorTrain x0 instead; the call is otherwise identical.

Precision — jit runs in jax’s default float32. use_jit=True opts into jax world, whose default dtype is float32. Fit a small exact-rank tensor three ways; the giveaway throughout is the returned core’s dtype. First set up the problem:

>>> import t3toolbox.manifold as t3m
>>> np.random.seed(0)
>>> A = t3.TuckerTensorTrain.randn((4, 5, 6), (2, 2, 2), (1, 2, 2, 1))
>>> ww = [np.random.randn(60, N) for N in (4, 5, 6)]
>>> ww = [w / np.linalg.norm(w, axis=1, keepdims=True) for w in ww]
>>> b = A.probe(ww)
>>> x0 = t3.TuckerTensorTrain.zeros((4, 5, 6), (2, 2, 2), (1, 2, 2, 1))
>>> def rel_err(x): return float(np.linalg.norm(np.asarray(x.to_dense()) - A.to_dense()) / np.linalg.norm(A.to_dense()))

  1. The numpy path is float64 and recovers to near machine precision (~1e-10):

>>> x_np, _ = optimizers.newton_cg(t3m.MANIFOLD, 'probe', ww, b, x0, max_newton=20)
>>> str(x_np.data[0][0].dtype), bool(rel_err(x_np) < 1e-8)
('float64', True)

  1. use_jit=True runs the SAME code in jax’s default float32 and stalls ~1000x coarser (~1e-7):

>>> x_jit, _ = optimizers.newton_cg(t3m.MANIFOLD, 'probe', ww, b, x0, max_newton=20, use_jit=True)
>>> str(x_jit.data[0][0].dtype), bool(rel_err(x_jit) > 1e-8)
('float32', True)

(3) Enabling jax x64 restores float64 under jit – full accuracy again (~1e-10). jax_enable_x64 is a global, process-wide flag, so restore it right after the fit; the check reads values captured while it was on:

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> x_x64, _ = optimizers.newton_cg(t3m.MANIFOLD, 'probe', ww, b, x0, max_newton=20, use_jit=True)
>>> dtype_x64, ok_x64 = str(x_x64.data[0][0].dtype), bool(rel_err(x_x64) < 1e-8)
>>> jax.config.update("jax_enable_x64", False)          # restore the default before asserting (no leak)
>>> dtype_x64, ok_x64
('float64', True)
Parameters:

  • kind (str)

  • sample (Any)

  • data (Any)

  • x0 (Point)

  • order (Optional[int])

  • weight (Optional[Any])

  • verbose (bool)

  • val_sample (Any)

  • val_data (Any)

  • callback (Optional[Callable])

  • use_jit (bool)

  • regularizer (Any)

  • chunk_size (Any)

Return type:

Tuple[Point, dict]