compile_optimizer

→_CompiledStepBase

compile_optimizer(opt: Optimizer)

source edit

Wrap opt so opt.step() runs as a single MPSGraph executable.

Dispatches on the concrete optimizer subclass and returns the matching _CompiledStepBase subclass instance. The returned object exposes the same eager-optimizer API surface (step / zero_grad / param_groups / state_dict / load_state_dict) so existing training loops slot it in without code changes — the only difference is that step() runs as one compiled GPU kernel instead of N sequential per- parameter element-wise updates.

Supported optimizers (8): * lucid.optim.SGD — classical, with optional momentum / Nesterov / weight-decay branches. * lucid.optim.Adam, lucid.optim.AdamW — bias-corrected first + second moments; coupled vs decoupled weight decay respectively. AMSGrad variant of Adam is rejected at construct time (a planned follow-up). * lucid.optim.RMSprop — exponentially-smoothed squared gradient. centered=True is rejected loudly here even though the eager backend silently drops it. * lucid.optim.Adagrad — per-parameter cumulative squared gradient; lr_decay is fed as a per-step scalar so the trace stays signature-stable. * lucid.optim.Adadelta — running RMS-ratio adaptive step (no manual LR). * lucid.optim.Adamax — Adam variant with L∞-norm second-moment estimate. * lucid.optim.NAdam — Adam with Nesterov lookahead + momentum-decay schedule (closed-form μ_t series).

Structurally unsupported (raise NotImplementedError with the reason at construct time): * lucid.optim.LBFGS — line search depends on data-dependent iteration count (no static-shape MPSGraph equivalent). * lucid.optim.SparseAdam — index-driven update requires runtime nonzero / scatter. * lucid.optim.Rprop — sign-based per-element conditional branching. * lucid.optim.RAdam — ρ_t > 4 branch can't be a static MPSGraph op. * lucid.optim.ASGD — averaging coefficient depends on iteration count past a warmup threshold.

Multi-param-group optimizers are not yet supported (a future pass will emit one update kernel per group); single-group optimizers cover essentially every practical training recipe.

Parameters

optOptimizer

Concrete optimizer instance whose step() is being lifted. Held by reference; LR-scheduler callbacks and other state mutations on opt continue to take effect because the compiled subclass reads them at scalar-refresh time, not at construct time.

Returns

_CompiledStepBase

Drop-in optimizer replacement. The underlying class is one of the _Compiled<Name> subclasses listed above.

Raises

NotImplementedError

With a structural reason when opt is one of the unsupported families above (or AMSGrad / multi-group / etc.).

TypeError

When opt is not a Lucid Optimizer subclass at all — the message lists the supported set.

Examples

Drop-in replacement:

    opt = lucid.optim.Adam(model.parameters(), lr=1e-3)
    copt = compile_optimizer(opt)

    for batch, target in loader:
        copt.zero_grad()
        loss = F.cross_entropy(model(batch), target)
        loss.backward()
        copt.step()             # one MPSGraph executable

Inspecting the rejection reason:

    try:
        copt = compile_optimizer(lucid.optim.LBFGS(params))
    except NotImplementedError as e:
        print(e)
        # "compile_optimizer: LBFGS is not supported.  LBFGS performs
        #  a line search inside step() whose iteration count depends
        #  on tensor values; ..."

Used by 2

Drop-in replacement: opt = lucid.optim.Adam(model.parameters(), lr=1e-3) copt = compile_optimizer(opt) for batch, target in loader: copt.zero_grad() loss = F.cross_entropy(model(batch), target) loss.backward() copt.step() # one MPSGraph executable Inspecting the rejection reason: try: copt = compile_optimizer(lucid.optim.LBFGS(params)) except NotImplementedError as e: print(e) # "compile_optimizer: LBFGS is not supported. LBFGS performs # a line search inside step() whose iteration count depends # on tensor values; ..."