autocast
autocast(device_type: str = 'metal', dtype: dtype = float16, enabled: bool = True)Enable automatic mixed-precision computation inside a context.
Operations that opt into AMP (matmul, conv, attention, …) cast their inputs to the target lower-precision dtype on entry, reducing memory pressure and — on Metal — improving throughput by activating the M-series GPU's half-precision tensor pipelines. Cast points are inserted only at op boundaries; activations themselves are stored in the lower-precision dtype but accumulator and reduction state stay in float32 to preserve numerical fidelity (the standard "mixed" recipe).
The engine's AutocastGuard is enter-only (its
__exit__ is a no-op), so this Python class implements full RAII
on top of it: __enter__ snapshots the previously active AMP
state, __exit__ reinstalls it (or installs a neutral float32
guard if AMP was inactive before the scope).
Parameters
device_typestr= 'metal'"metal" (GPU stream,
default) or "cpu" (Accelerate stream).dtypelucid.dtype= float16lucid.float16.enabledbool= TrueFalse the context is a no-op — useful for ablation /
A/B comparing AMP on vs off without changing call sites.
Default True.Examples
Context-manager form (typical training loop):
>>> import lucid
>>> with lucid.amp.autocast():
... output = model(input) # ops cast to float16 inside
... loss = criterion(output, target)
Decorator form (every call wraps itself in a fresh scope):
>>> @lucid.amp.autocast()
... def predict(x):
... return model(x)
Nested scopes restore the outer dtype on exit:
>>> with lucid.amp.autocast(dtype=lucid.float16):
... ... # fp16 here
... with lucid.amp.autocast(dtype=lucid.bfloat16):
... ... # bf16 here
... ... # back to fp16 — prior guard reinstalledSee Also
- lucid.amp.GradScaler—loss-scaling counterpart for fp16 training.
Used by 1
Constructors
2__init__
→None__init__(device_type: str = 'metal', dtype: dtype = float16, enabled: bool = True)Configure the autocast context.
Parameters
device_typestr= 'metal''metal' (GPU stream) and 'cpu' (Accelerate stream).dtypelucid.dtype= float16enabledbool= TrueFalse the context is a no-op (useful for ablation).Use as a function decorator.
Static methods
2Dunder methods
2Activate the autocast scope and return self.
Captures the previously active AMP state (if any) so it can be
restored on exit, then installs a new AutocastGuard for the
configured target dtype.
Restore the previous AMP state on scope exit.
If autocast was already active before this scope, the prior
dtype guard is reinstalled. Otherwise a neutral float32 guard
is entered since the engine has no disable_amp() primitive.