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. 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 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
>>> 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 = 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.