lgatr.utils.autocast

Pin inputs to a minimum autocast precision; usable as a decorator.

Functions

autocast_dtype([device_type])

Dtype that autocast would cast to on device_type.

Classes

`minimum_autocast_precision`([min_dtype, output])	Pin tensors to a minimum precision inside autocast regions.
`naive_amp`([enabled])	Disable all `minimum_autocast_precision` pinning inside the block.

lgatr.utils.autocast.autocast_dtype(device_type='cuda')[source]

Dtype that autocast would cast to on device_type.

Return type:: dtype

class lgatr.utils.autocast.minimum_autocast_precision(min_dtype=torch.float32, output='low')[source]

Pin tensors to a minimum precision inside autocast regions.

Used as a decorator: @minimum_autocast_precision(torch.float32) on a function definition. Inside autocast-enabled regions, floating-point inputs below min_dtype are cast up to min_dtype, autocast is disabled for the call, and outputs are optionally cast per the output argument. Outside autocast regions the decorator is a no-op.

The naive_amp context manager turns the decorator into a no-op, letting the wrapped ops run in the surrounding autocast dtype instead.

Only floating-point tensors are modified — non-tensors, integer tensors, and boolean tensors are left alone.

Note: AMP is enabled separately for CPU and CUDA. This may behave unexpectedly when both devices are used and only one of them has AMP enabled. Instances are not thread-safe; share one per thread if used concurrently.

Parameters:

min_dtype (dtype) – Minimum dtype.
output (Union[Literal['low', 'high'], dtype, None]) – Specifies which dtype the outputs should be cast to. Only floating-point tensor outputs are affected. If "low" (default), the lowest precision among min_dtype and the input dtypes is used. If "high", the highest-precision input dtype is used. If None, outputs are not modified. If a torch.dtype, that dtype is used. In the "low" and "high" modes, outputs are left alone when there are no floating-point inputs to derive a dtype from.

cast(var)[source]

Upcast a floating-point tensor to at least min_dtype.

Return type:: Any

class lgatr.utils.autocast.naive_amp(enabled=True)[source]

Disable all minimum_autocast_precision pinning inside the block.

While active, the fp32 precision islands created by the minimum_autocast_precision decorator are bypassed and the wrapped ops run in the surrounding autocast dtype (e.g. bf16). Restores the previous state on exit; safe to nest.

Parameters:: enabled (bool) – Whether to enable naive-AMP mode. False leaves the current state untouched, making naive_amp(False) a no-op that still nests cleanly (it never overrides an outer naive_amp).