adam

t3toolbox.backend.optimizers.adam(problem, x0, rng, batch, draw=None, lr=0.01, max_iter=2000, betas=(0.9, 0.999), eps=1e-08, cosine=True, use_jit=False)

def adam(
        problem:  Problem,    # the fixed-rank least-squares problem
        x0:       Tangent,    # initial cores (U, G)
        rng,                  # np.random.Generator -- passed to the draw each step
        batch:    int,        # measurements per minibatch (for the default flat draw; ignored if draw given)
        draw:     typ.Optional[typ.Callable] = None,  # custom draw(rng)->(sample_B,data_B); None = flat default
        lr:       float = 1e-2,
        max_iter: int   = 2000,
        betas:    typ.Tuple[float, float] = (0.9, 0.999),
        eps:      float = 1e-8,
        cosine:   bool  = True,   # cosine-decay the learning rate over max_iter (helps it settle)
        use_jit:  bool  = False,  # jit the per-step kernel (gradient + Adam update + retract) when jax
) -> typ.Tuple[Tangent, dict]:    # (x_cores, stats)

Adam over the cores – the dependency-free first-order method for the corewise geometry. The moments m / v are trees matching the cores (a corewise_map per step); elementwise, so this is exactly per-core Adam. (lr is a real hyperparameter – Adam is not tuning-free, unlike the Cauchy step.) On corewise the step is additive (lm.retract = cores += step). use_jit jits the per-step kernel (lr_t/t flow in as traced args, so the schedule does not force a recompile), auto-converting numpy inputs to jax first (_prepare_jit_inputs()).

Parameters:

• problem (Problem)

• x0 (Tangent)

• batch (int)

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

• lr (float)

• max_iter (int)

• betas (t3toolbox.backend.common.typ.Tuple[float, float])

• eps (float)

• cosine (bool)

• use_jit (bool)

Return type:

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