Executable

class nvtripy.Executable(executable, arg_names, output_devices)

Bases: object

Represents a compiled executable generated by the compiler.

See also

compile()

load() → Executable

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/tmpsymvtzt7 does not exist, creating now.

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

Invokes the executable with the specified tensor arguments.

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)
b = tp.ones((1,), dtype=tp.float32)

out = compiled_add(a, b)

Local Variables

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

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

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

save(path: str) → None

Saves this executable to the provided path.

Parameters:

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

Return type:

None

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/tmpi26binzu does not exist, creating now.