lgatr.primitives.compile

Helpers for using L-GATr primitives with torch.compile().

Functions

warmup_caches(device[, dtype])

Pre-populate L-GATr's primitive caches for a given device and dtype.

lgatr.primitives.compile.warmup_caches(device, dtype=torch.float32)[source]

Pre-populate L-GATr’s primitive caches for a given device and dtype.

Without warming, the first call for a new (device, dtype) pair triggers host-to-device copies that partition the captured graph under torch.compile() with mode="reduce-overhead". Calling this helper once per (device, dtype) before compiling avoids the partition.

Parameters:

device (device | str) – Target device, either a torch.device or a string like "cuda" or "cpu".
dtype (dtype) – Floating-point dtype matching what the model will run in.

Return type:

None