Executable¶

class nvtripy.Executable[source]¶

Bases: object

Represents a compiled executable generated by the compiler.

See also

compile()

input_infos: Dict[str, InputInfo]¶: Stores metadata, like shapes and data types, for each input to the executable.

load() → Executable[source]¶

Loads a executable from the provided path.

Parameters:: path (str) – The path from which to load the exectuable.
Returns:: The executable object loaded from the file.
Return type:: Executable

Example: Save and load executable

import os


def add(a, b):
    return a + b


compiled_add = tp.compile(
    add,
    args=[
        tp.InputInfo(shape=((1, 2, 3),), dtype=tp.float32),
        tp.InputInfo(shape=((1, 2, 3),), dtype=tp.float32),
    ],
)


# Assuming `out_dir` is the directory containing the executable:
executable_file = os.path.join(out_dir, "executable.json")
loaded_executable = tp.Executable.load(executable_file)

Local Variables¶

>>> loaded_executable
Executable(a: nvtripy.Tensor, b: nvtripy.Tensor) -> nvtripy.Tensor

Output¶

[I] /tmp/tmpj1_xhnt4 does not exist, creating now.

__call__(*args: Tensor, **kwargs: Tensor) → Tensor | Sequence[Tensor][source]¶

Invokes the executable with the specified tensor arguments.

Note

Inputs must be evaluated tensors in GPU memory.

You can use nvtripy.copy() or nvtripy.Tensor.eval() to ensure this.

Parameters:

*args (Tensor) – Positional arguments. Must be of type Tensor .
**kwargs (Tensor) – Keyword arguments. Must be of type Tensor .

Returns:

The output Tensor s of the compiled function.

Return type:

Tensor | Sequence[Tensor]

Example

def add(a, b):
    return a + b


compiled_add = tp.compile(
    add,
    args=[
        tp.InputInfo((1,), dtype=tp.float32),
        tp.InputInfo((1,), dtype=tp.float32),
    ],
)

a = tp.ones((1,), dtype=tp.float32).eval()
b = tp.ones((1,), dtype=tp.float32).eval()

out = compiled_add(a, b)

Local Variables¶

>>> a
tensor([1], dtype=float32, loc=gpu:0, shape=(1,))

>>> b
tensor([1], dtype=float32, loc=gpu:0, shape=(1,))

>>> out
tensor([2], dtype=float32, loc=gpu:0, shape=(1,))

save(path: str) → None[source]¶

Saves this executable to the provided path.

Parameters:: path (str) – The path at which to save the executable.

Example: Save executable

import os


def add(a, b):
    return a + b


compiled_add = tp.compile(
    add,
    args=[
        tp.InputInfo(shape=((1, 2, 3),), dtype=tp.float32),
        tp.InputInfo(shape=((1, 2, 3),), dtype=tp.float32),
    ],
)

# Assuming `out_dir` is the desired output directory:
executable_file = os.path.join(out_dir, "executable.json")
compiled_add.save(executable_file)

Output¶

[I] /tmp/tmps2n56d3p does not exist, creating now.

property serialized_tensorrt_engine: bytes¶

The serialized TensorRT engine, as bytes, from the executable.