fitting#

Frontend for the least-squares fitting layer: one geometry-generic Gauss-Newton model.

A GaussNewtonModel is the local Gauss-Newton model of a least-squares objective ½‖F(X) − y‖² at a fixed point, linearized through a geometry: it exposes the objective value c = ½‖r‖², the gradient g, the Gauss-Newton Hessian action H p = JᵀJ p, and the quadratic-model value m(p) = c + gᵀp + ½ pᵀ H p. The model is generic over two independent choices:

the geometry – MANIFOLD (optimize on the fixed-rank manifold; the gauge projection Π makes g / H Riemannian) or COREWISE (optimize the raw cores; no Π). The geometry supplies the frame frame = geometry.frame(X) and the projection Π = geometry.project; the forward is J = 𝒥∘Π and the gradient is Jᵀr = Π∘𝒥ᵀr. Manifold ⟺ Π, corewise ⟺ no Π is structural (bundled in the geometry), never a flag – mixing the two silently corrupts the result.

the sampling kind – bound by the factory apply_model() / entries_model() / probe_model() (the bare 𝒥 / 𝒥ᵀ from t3toolbox.backend.probing, bundled in t3toolbox.backend.fitting).

Every input and output is a T3Tangent at model.frame – including the corewise case, where the tangent lives at the non-orthonormal frame (U,G,G,G) (a core perturbation; see CorewiseGeometry). Build trial steps at model.frame (e.g. geometry.randn(model.frame)); a step at a different frame raises the numerical same-frame guard (skipped under safety.unsafe() / a jax trace), like T3Tangent.

The frame sweep (the frame-and-data edge variables) is computed once by the factory and stored as the sweep field, reused across every gradient / gn_hessian / evaluate – in an inner CG the frame is fixed, so the sweep is computed once, not once per matrix-vector product. See t3toolbox.backend.fitting and dev/archive/geometry_refactor_plan.md.

Jitting an optimizer#

GaussNewtonModel is a registered jax pytree (the data – frame, sweep, sample, residual – are leaves; geometry / kind are static aux). Crucially the frame flows as a leaf, not aux (the same is true of T3Tangent’s frame), so a model or tangent that crosses a jit boundary does NOT recompile when the frame changes – the per-frame recompile that frame-as-aux used to force is gone. So you can jit the frontend matvec directly:

Inner-solve jit (Newton-CG). Jit the matvec with the model and the tangent as arguments; it compiles once for the whole solve (the model’s frame/sweep change as data, not as a recompile trigger), reused across every outer step and CG iteration:
```
Hmatvec = jax.jit(lambda model, p: model.gn_hessian(p))
# per outer step:  model = fitting.apply_model(geom, X, ww, r)
# inner CG:         Hmatvec(model, p_k)
```

Whole-step jit (Cauchy / gradient descent / a fixed-step loop). Jit a step function whose only argument is the point X and build the model inside; compiles once, reused across steps:

@jax.jit
def step(X):
    r     = X.apply(ww) - b
    model = fitting.apply_model(t3m.MANIFOLD, X, ww, r)
    g     = model.gradient
    alpha = g.corewise_inner(g) / model.gn_quadratic(g)     # cheap Cauchy step (one forward)
    return t3m.MANIFOLD.retract((-alpha) * g)

Under a trace the numerical same-frame guards skip (you cannot branch on a tracer; jit is unsafe mode), so a matvec output combines with the eager CG iterates freely. The geometry singletons (t3m.MANIFOLD / COREWISE) are zero-leaf pytrees – close over or pass as args freely.

Classes#

`GaussNewtonModel`	The local Gauss-Newton model of a least-squares objective at `frame`, generic over the geometry.
`UniformGaussNewtonModel`	The uniform-layer twin of `GaussNewtonModel`: the local Gauss-Newton model at `frame`,

Functions#

`apply_model`(geometry, x, ww, residual[, regularizer])	The Gauss-Newton model of an all-modes `apply` least-squares objective at `x`, on `geometry`.
`entries_model`(geometry, x, index, residual[, regularizer])	The Gauss-Newton model of an all-modes `entries` least-squares objective at `x`, on `geometry`.
`probe_model`(geometry, x, ww, residual[, weight, ...])	The Gauss-Newton model of a `probe` least-squares objective at `x`, on `geometry`.
`apply_derivatives_model`(geometry, x, ww, pp, order, ...)	The Gauss-Newton model of an `apply`-derivatives least-squares objective at `x`: the
`entries_derivatives_model`(geometry, x, index, pp, ...)	The `entries`-derivatives Gauss-Newton model -- like `apply_derivatives_model()` at integer
`probe_derivatives_model`(geometry, x, ww, pp, order, ...)	The `probe`-derivatives Gauss-Newton model -- vector-valued (one free mode per probe), so