torch_staintools.functional.compile package

Module contents

class torch_staintools.functional.compile.CompiledWrapper(*args, **kwargs)

Bases: Protocol

compiled_fn: Callable | None

reset_cache() → None

torch_staintools.functional.compile.lazy_compile(func: Callable) → CompiledWrapper

Enable or disable torch.compile by torch_staintools.constants.CONFIG.ENABLE_COMPILE.

If True, function will be compiled and cached. Otherwise, it will be executed in eager mode.

Parameters:

func – The function to compile.

Returns:

The compiled function or the original function.

Return type:

CompiledWrapper

torch_staintools.functional.compile.static_compile(func: Callable) → Callable

Import-time wrapper.

CONFIG.ENABLE_COMPILE must be modified before importing any compiled functions.

Parameters:

func – The function to compile.

Returns:

The compiled function or the original function.