compile

torch_tvarant.compiler.compile(module, *, dynamic=False) → Callable

Compile a module for Tvarant inference.

Prefers FX tracing that inlines Linear / activations (stable on this custom device). Falls back to torch.compile(..., backend="tvarant").

Parameters:

Name	Type	Description	Default
module	Module	Module to optimize. Switched to eval() mode.	required
dynamic	bool	Forwarded to torch.compile on the fallback path only.	False

Returns:

Name	Type	Description
Callable	Callable	Callable module / optimized callable producing the same
	Callable	outputs as module on device="tvarant".

Examples:

>>> compiled = torch_tvarant.compiler.compile(model)
>>> y = compiled(x)

Source code in torch_tvarant/compiler.py

def compile(module: torch.nn.Module, *, dynamic: bool = False) -> Callable:
    """Compile a module for Tvarant inference.

    Prefers FX tracing that inlines Linear / activations (stable on this custom
    device). Falls back to ``torch.compile(..., backend="tvarant")``.

    Args:
        module: Module to optimize. Switched to ``eval()`` mode.
        dynamic: Forwarded to ``torch.compile`` on the fallback path only.

    Returns:
        Callable: Callable module / optimized callable producing the same
        outputs as ``module`` on ``device="tvarant"``.

    Examples:
        >>> compiled = torch_tvarant.compiler.compile(model)
        >>> y = compiled(x)
    """
    module.eval()
    try:
        return compile_fx(trace_module(module))
    except Exception:
        return torch.compile(module, backend="tvarant", dynamic=dynamic)