adam
====

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

   .. code-block:: python

      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 (:py:func:`_prepare_jit_inputs`).