optimizers#

Backend-first optimization algorithms for fixed-rank T3 least-squares fitting.

These operate on raw cores / tangent tuples via backend functions only – no frontend imports (no TuckerTensorTrain / T3Tangent / GaussNewtonModel / geometry classes). A raw-.data user can call them directly; the frontend optimizers.py is a thin adapter that validates the input once and assembles the Problem oracle from the same backend functions (design: dev/archive/optimizers_plan.md).

Because the numerical safety preconditions live only in the frontend, this layer is check-free of them (structural shape guards inside the backend functions remain, and are jit-safe). So jit needs no unsafe() wrapping; the use_jit machinery is a thin jax-only layer on top of the ordinary numpy/jax dispatch (the per-step kernels jit via _maybe_jit; the inner CG runs as a lax.while_loop – eager-vs-jit agreement is covered in tests/backend/test_optimizers.py).

The oracle:

  • Problem – (geometry ops, sampling kind, sample, data); builds per-point LocalModel objects and retracts.

  • LocalModel – the GN model linearized at a point (.gradient / .objective / .hvp / .gn_quadratic / .retract), the backend twin of fitting.GaussNewtonModel.

  • GeometryOps – (frame, project, retract) on raw data; MANIFOLD_OPS / COREWISE_OPS singletons.

Tangent vectors are raw (tucker_var, tt_var) tuples; vector arithmetic is the corewise ops.

Attributes#

COREWISE_OPS

MANIFOLD_OPS

Classes#

GeometryOps

LocalModel

The Gauss-Newton model linearized at a point -- the backend twin of fitting.GaussNewtonModel.

Problem

A fixed-rank least-squares fitting problem min_x ½‖S(x) - data‖² for a sampling op S on

NewtonInfo

Per-iteration diagnostics passed to a newton_cg callback -- everything one Newton step

Functions#

least_squares_problem(geom, kind, sample, data[, ...])

Assemble a least-squares Problem from a geometry, a sampling kind, the sample, and the observed

flat_draw(problem, batch)

The default minibatch draw: a uniform random subset of batch measurements across the whole

gradient_descent(problem, x0[, n_iter, gtol_rel, c_armijo])

Steepest descent with a Cauchy initial step and an Armijo backtracking line search. The step

mc_sgd(problem, x0, rng, batch[, draw, max_iter, ...])

Manifold Cauchy SGD (T4S §5.3.2): minibatch + the tuning-free Cauchy step

adam(problem, x0, rng, batch[, draw, lr, max_iter, ...])

Adam over the cores -- the dependency-free first-order method for the corewise geometry. The

newton_cg(problem, x0[, max_newton, gtol_rel, ...])

Inexact Riemannian Newton-CG with an Armijo line search -- the manifold workhorse (the gauged H

Module Contents#

COREWISE_OPS#
MANIFOLD_OPS#