Skip to content

Getting started

After installing torch_tvarant, verify the device and run a small model.

Verify installation

import torch
import torch_tvarant

print(torch.tvarant.is_available())   # True
print(torch.tvarant.backend())        # 'sim' or 'opencl'
x = torch.ones(4, device="tvarant")
print(x.device)                       # tvarant:0

Run a small MLP

import torch
import torch.nn as nn
import torch_tvarant

model = nn.Sequential(
    nn.Linear(16, 32),
    nn.ReLU(),
    nn.Linear(32, 4),
).to("tvarant")

x = torch.randn(2, 16, device="tvarant")
y = model(x)
print(y.shape)  # torch.Size([2, 4])

Compile for inference

import torch_tvarant

compiled = torch_tvarant.compiler.compile(model)
y = compiled(x)

See JIT Compiler for fusion details.

Learn more

Topic Page
Device helpers torch.tvarant
Supported ops Supported ops
Kernels Kernels
C++ extension C++ API
FPGA path FPGA / OpenCL