lloca.utils.compile

Helpers for using lloca networks with torch.compile().

Functions

compile_model(model, *[, compile_kwargs])

Wrap model.forward with torch.compile() in place.
lloca.utils.compile.compile_model(model, *, compile_kwargs=None)[source]

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.

Return type:

None