aitune.torch.backend.torch_tensorrt_jit_backend

Torch compile backend.

Module Contents

Classes

Name	Description
`TorchTensorRTConfig`	Fallback settings matching upstream defaults when torch_tensorrt is unavailable.
`TorchTensorRTJitBackend`	Torch TensorRT thought torch.compile(backend=“torch_tensorrt”).
`TorchTensorRTJitBackendConfig`	Configuration for torch.compile(backend=“torch_tensorrt”).

Functions

Name	Description
`assert_torch_tensorrt`	Check if torch_tensorrt is installed.

Data

logger

API

class aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTConfig(
    enabled_precisions: set[torch.dtype] | None = None,
    use_explicit_typing: bool = True,
    workspace_size: int = 0
)

Dataclass

Fallback settings matching upstream defaults when torch_tensorrt is unavailable.

enabled_precisions

set[dtype] | None = None

use_explicit_typing

bool = True

workspace_size

int = 0

class aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend(
    config: aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackendConfig | None = None
)

Bases: Backend

Torch TensorRT thought torch.compile(backend=“torch_tensorrt”).

Backend does not use intermediate formats, and compiled model is not stored.

STATE_COMPILE_DYNAMIC

= 'compile_dynamic'

STATE_CONFIG

= 'config'

STATE_DATA

= 'data'

STATE_DEVICE

= 'device'

STATE_ORIG_MODULE

= 'orig_module'

STATE_TYPE

= 'type'

_compile_dynamic

= self._config.dynamic

_config

= config or TorchTensorRTJitBackendConfig()

_devices

list[str] = ['cuda']

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend._activate()

Activate backend.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend._build(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    cache_dir: pathlib.Path
) -> aitune.torch.backend.backend.Backend

Build the model with Torch compile.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend._compile()

Compile module with Torch compile.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend._deactivate()

Deactivate backend and cleanup.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend._deploy()

Deploys the backend.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend._infer(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Run inference with TensorRT engine thought Torch compile.

Parameters:

*args

AnyDefaults to ()

Inference arguments

**kwargs

AnyDefaults to {}

Inference keyword arguments

Returns: Any

Model outputs

Raises:

AITuneError: If backend is not activated

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend._save_config(
    cache_dir: pathlib.Path
)

Store the backend configuration to a file.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend.describe() -> str

Returns the description of the backend.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend.from_dict(
    module: torch.nn.Module | None,
    state_dict: dict
)

classmethod

Creates a backend from a state_dict.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend.key() -> str

Returns the key of the backend.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackend.to_dict()

Returns the state_dict of the backend.

class aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackendConfig(
    compile_config: aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTConfig = TorchTensorRTConfig(),
    fullgraph: bool = False,
    dynamic: bool | None = None,
    autocast_enabled: bool = False,
    autocast_dtype: torch.dtype | None = None
)

Dataclass

Bases: BackendConfig

Configuration for torch.compile(backend=“torch_tensorrt”).

Defaults preserve torch.compile behavior unless callers opt into stricter compilation options.

autocast_dtype

dtype | None = None

autocast_enabled

bool = False

compile_config

TorchTensorRTConfig = field(default_factory=TorchTensorRTConfig)

dynamic

bool | None = None

fullgraph

bool = False

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackendConfig.describe() -> str

Describe the backend configuration. Display only changed fields.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackendConfig.from_dict(
    data: dict
) -> aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackendConfig

classmethod

Initialise config from a plain dict (e.g. parsed from YAML).

compile_config may be passed as a nested dict and will be reconstructed into a TorchTensorRTConfig instance automatically.

aitune.torch.backend.torch_tensorrt_jit_backend.TorchTensorRTJitBackendConfig.to_dict()

Convert TorchTensorRTJitBackendConfig to dict.

aitune.torch.backend.torch_tensorrt_jit_backend.assert_torch_tensorrt()

Check if torch_tensorrt is installed.

aitune.torch.backend.torch_tensorrt_jit_backend.logger = getLogger(__name__)