JIT Configuration API

Config

aitune.torch.jit.config.Config `dataclass`

Config(dry_run=False, dry_run_failure_probability=0.2, inspect_mode=False, device=DEFAULT_DEVICE, min_samples=1, batch_axis_required=True, max_depth_level=1, min_parameters=0, detect_graph_breaks=False, skip_modules=list(), cache_dir=Path(get('AITUNE_JIT_CACHE_DIR', DEFAULT_JIT_CACHE_DIR)), backends=(lambda: [TensorRTBackend(config=TensorRTBackendConfig(use_dynamo=True)), TensorRTBackend(config=TensorRTBackendConfig(use_dynamo=False)), TorchInductorBackend()])())

Configuration for JIT module.

__post_init__

__post_init__()

Post init.

Source code in aitune/torch/jit/config.py

def __post_init__(self):
    """Post init."""
    self.device = get_device(self.device)

reset_to_defaults

reset_to_defaults()

Reset all options to their default values (e.g. for test isolation).

Source code in aitune/torch/jit/config.py

def reset_to_defaults(self) -> None:
    """Reset all options to their default values (e.g. for test isolation)."""
    defaults = Config()
    for f in fields(Config):
        setattr(self, f.name, getattr(defaults, f.name))

config

Global JIT configuration instance:

aitune.torch.jit.config