lgatr.utils.compile

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

Functions

compile_model(model, *[, compile_kwargs, ...])	Wrap model.forward with torch.compile() in place.
warmup_after_apply(model)	Warm L-GATr's primitive caches for the model's current device and dtype.

lgatr.utils.compile.compile_model(model, *, compile_kwargs=None, activation_memory_budget=None)

Wrap model.forward with torch.compile() in place.

Rebinding self.forward rather than patching the class keeps the compilation local to this instance.

Parameters:

• model (Module) – The torch.nn.Module whose forward should be compiled.

• compile_kwargs (Mapping | None) – Forwarded verbatim to torch.compile() (e.g. mode, dynamic, fullgraph, backend). Any key omitted falls back to torch’s own default.

• activation_memory_budget (float | None) – Fraction in [0, 1] for the partitioner’s activation-memory budget; lower values trade backward FLOPs for a smaller activation memory peak. None (default) leaves torch’s global setting untouched. Applied via a scoped patch only in effect while this model (re)compiles. Requires torch>=2.4.

Return type:

None

lgatr.utils.compile.warmup_after_apply(model)

Warm L-GATr’s primitive caches for the model’s current device and dtype.

Intended to be called from a torch.nn.Module._apply() override so the caches are populated whenever the model is moved or cast (.to() / .cuda() / .float() / etc.).

Parameters:

model (Module) – The torch.nn.Module whose primitive caches should be warmed.

Return type:

None