autocast
autocast(device_type: str = 'metal', dtype: dtype | None = None, 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).
Leaving the scope restores exactly the state it was entered from —
autocast off, or the enclosing scope's dtype — whether it is left
normally or by an exception, and however long the autocast object
itself lives on afterwards. One object can be entered again after it
exits, entered while it is already active (each entry is undone by its
own exit), and used as a decorator. enabled=False disables
autocast on device_type for the duration of the scope while leaving
the other device's autocast state untouched.
Autocast state is tracked independently per device ("metal" vs
"cpu") on each thread.
Parameters
device_typestr= 'metal'"metal" (GPU stream,
default) or "cpu" (Accelerate stream); anything else raises
ValueError.dtypelucid.dtype= Nonelucid.float16. On the CPU, float16
runs in float32 — Accelerate has no half-precision kernels — so a
CPU scope with the default dtype changes nothing; pass
dtype=lucid.bfloat16 for reduced precision there, as the
reference framework's CPU autocast does by default. Not every CPU
op has a bfloat16 kernel yet.enabledbool= TrueFalse autocast is disabled on device_type inside the
scope — useful for ablation or forcing full precision in a nested
subregion. Default True.Examples
Context-manager form (typical training loop):
>>> import lucid
>>> import lucid.nn as nn
>>> model = nn.Linear(4, 2).to("metal")
>>> x = lucid.randn(8, 4, device="metal")
>>> with lucid.amp.autocast():
... output = model(x) # ops cast to float16 inside
>>> output.dtype
lucid.float16
Decorator form (every call wraps itself in a fresh scope):
>>> @lucid.amp.autocast()
... def predict(x):
... return model(x)
>>> predict(x).dtype
lucid.float16
Nested scopes restore the outer dtype on exit:
>>> a = lucid.randn(2, 2, device="metal")
>>> with lucid.amp.autocast(dtype=lucid.float16):
... print((a @ a).dtype) # fp16 here
... with lucid.amp.autocast(dtype=lucid.bfloat16):
... print((a @ a).dtype) # bf16 here
... print((a @ a).dtype) # back to fp16 — prior guard reinstalled
lucid.float16
lucid.bfloat16
lucid.float16See Also
- lucid.amp.GradScaler—loss-scaling counterpart for fp16 training.
Used by 1
Constructors
2__init__
→None__init__(device_type: str = 'metal', dtype: dtype | None = None, enabled: bool = True)Configure the autocast context.
Parameters
device_typestr= 'metal''metal' (GPU stream) and 'cpu' (Accelerate stream).dtypelucid.dtype or None= Nonelucid.float16 when None).enabledbool= TrueFalse autocast is disabled on device_type inside
the scope (useful for ablation or nested full-precision blocks).Raises
ValueErrordevice_type is neither 'metal' nor 'cpu'.Use as a function decorator.
Every call of the decorated function runs inside this scope — the
same device_type, dtype and enabled — and leaves it on
return or on an exception.
Parameters
fncallableReturns
callablefn wrapped so that each call enters and exits the scope.
Static methods
2Return the currently active autocast dtype (for device_type if given), or None.
Return True if autocast is currently active (for device_type if given).
Dunder methods
2Activate (or disable) the autocast scope and return self.
Constructing the engine AutocastGuard sets the AMP state for
self._device_type to (enabled, dtype) and records the state
it replaced; the guard is kept until __exit__.
Restore the AMP state that was active when the scope was entered.
Releases the guard __enter__ made. Its destructor runs at
once and puts back the state it recorded, so the scope ends here
rather than whenever the autocast object is collected.
Exceptions propagate.