wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
apache--tvm/tests/python/relax/test_frontend_from_fx.py
T
wehub-resource-sync 26446540fa
Lint / lint (push) Waiting to run
Details
CI / MacOS (push) Waiting to run
Details
CI / Windows (push) Waiting to run
Details
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

6829 lines
227 KiB
Python
Raw Permalink Blame History

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: F841
import math
import operator
import pytest
import torch
import torch.nn.functional as F
from torch import fx
from torch.nn import Module
import tvm
import tvm.testing
from tvm import relax
from tvm.relax.frontend import detach_params
from tvm.relax.frontend.torch import from_fx
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.testing import env
def verify_model(torch_model, input_info, binding, expected):
graph_model = fx.symbolic_trace(torch_model)
with torch.no_grad():
mod = from_fx(graph_model, input_info)
binding = {k: tvm.runtime.tensor(v) for k, v in binding.items()}
expected = relax.transform.BindParams("main", binding)(expected)
tvm.ir.assert_structural_equal(mod, expected)
def test_conv1d():
class Conv1D1(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv1d(3, 6, 7, bias=True)
def forward(self, input):
return self.conv(input)
class Conv1D1Func(Module):
def __init__(self):
super().__init__()
self.weight = torch.randn(size=[6, 3, 7])
self.bias = torch.randn(size=[6])
def forward(self, input):
return torch.nn.functional.conv1d(input, self.weight, self.bias)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10), dtype="float32"),
w1: R.Tensor((6, 3, 7), dtype="float32"),
w2: R.Tensor((6,), dtype="float32"),
) -> R.Tensor((1, 6, 4), dtype="float32"):
# block 0
with R.dataflow():
lv1: R.Tensor((1, 6, 4), dtype="float32") = R.nn.conv1d(
input_1,
w1,
strides=[1],
padding=[0, 0],
dilation=[1],
data_layout="NCW",
kernel_layout="OIW",
out_layout="NCW",
out_dtype="float32",
)
lv2: R.Tensor((1, 6, 1), dtype="float32") = R.reshape(w2, [1, 6, 1])
lv3: R.Tensor((1, 6, 4), dtype="float32") = R.add(lv1, lv2)
gv: R.Tensor((1, 6, 4), dtype="float32") = lv3
R.output(gv)
return gv
class Conv1D2(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv1d(3, 6, 7, bias=False)
def forward(self, input):
return self.conv(input)
@tvm.script.ir_module
class expected2:
@R.function
def main(
input_1: R.Tensor((1, 3, 10), dtype="float32"),
w1: R.Tensor((6, 3, 7), dtype="float32"),
) -> R.Tensor((1, 6, 4), dtype="float32"):
# block 0
with R.dataflow():
lv1: R.Tensor((1, 6, 4), dtype="float32") = R.nn.conv1d(
input_1,
w1,
strides=[1],
padding=[0, 0],
dilation=[1],
data_layout="NCW",
kernel_layout="OIW",
out_layout="NCW",
out_dtype="float32",
)
gv: R.Tensor((1, 6, 4), dtype="float32") = lv1
R.output(gv)
return gv
input_info = [([1, 3, 10], "float32")]
model = Conv1D1()
binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
verify_model(model, input_info, binding, expected1)
model = Conv1D1Func()
binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
verify_model(model, input_info, binding, expected1)
model = Conv1D2()
binding = {"w1": model.conv.weight.detach().numpy()}
verify_model(model, input_info, binding, expected2)
def test_conv1d_transpose():
class ConvTranspose1d1(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.ConvTranspose1d(6, 6, 3, bias=True)
def forward(self, input):
return self.conv(input)
class ConvTranspose1d1Func(Module):
def __init__(self):
super().__init__()
self.weight = torch.randn(size=[6, 6, 3])
self.bias = torch.randn(size=[6])
def forward(self, input):
return torch.nn.functional.conv_transpose1d(input, self.weight, self.bias)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 6, 4), dtype="float32"),
w1: R.Tensor((6, 6, 3), dtype="float32"),
w2: R.Tensor((6,), dtype="float32"),
) -> R.Tensor((1, 6, 6), dtype="float32"):
# block 0
with R.dataflow():
lv1: R.Tensor((1, 6, 6), dtype="float32") = R.nn.conv1d_transpose(
input_1,
w1,
strides=[1],
padding=[0, 0],
output_padding=[0],
dilation=[1],
data_layout="NCW",
kernel_layout="IOW",
out_layout="NCW",
out_dtype="float32",
)
lv2: R.Tensor((1, 6, 1), dtype="float32") = R.reshape(w2, [1, 6, 1])
lv3: R.Tensor((1, 6, 6), dtype="float32") = R.add(lv1, lv2)
gv: R.Tensor((1, 6, 6), dtype="float32") = lv3
R.output(gv)
return gv
class ConvTranspose1d2(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.ConvTranspose1d(6, 6, 3, bias=False)
def forward(self, input):
return self.conv(input)
@tvm.script.ir_module
class expected2:
@R.function
def main(
input_1: R.Tensor((1, 6, 4), dtype="float32"),
w1: R.Tensor((6, 6, 3), dtype="float32"),
) -> R.Tensor((1, 6, 6), dtype="float32"):
# block 0
with R.dataflow():
lv1: R.Tensor((1, 6, 6), dtype="float32") = R.nn.conv1d_transpose(
input_1,
w1,
strides=[1],
padding=[0, 0],
output_padding=[0],
dilation=[1],
data_layout="NCW",
kernel_layout="IOW",
out_layout="NCW",
out_dtype="float32",
)
gv: R.Tensor((1, 6, 6), dtype="float32") = lv1
R.output(gv)
return gv
input_info = [([1, 6, 4], "float32")]
model = ConvTranspose1d1()
binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
verify_model(model, input_info, binding, expected1)
model = ConvTranspose1d1Func()
binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
verify_model(model, input_info, binding, expected1)
model = ConvTranspose1d2()
binding = {"w1": model.conv.weight.detach().numpy()}
verify_model(model, input_info, binding, expected2)
def test_conv2d():
class Conv2D1(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(3, 6, 7, bias=True)
def forward(self, input):
return self.conv(input)
class Conv2D1Func(Module):
def __init__(self):
super().__init__()
self.weight = torch.randn(size=[6, 3, 7, 7])
self.bias = torch.randn(size=[6])
def forward(self, input):
return torch.nn.functional.conv2d(input, self.weight, self.bias)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((6, 3, 7, 7), dtype="float32"),
w2: R.Tensor((6,), dtype="float32"),
) -> R.Tensor((1, 6, 4, 4), dtype="float32"):
# block 0
with R.dataflow():
lv1: R.Tensor((1, 6, 4, 4), dtype="float32") = R.nn.conv2d(
input_1,
w1,
strides=[1, 1],
padding=[0, 0, 0, 0],
dilation=[1, 1],
data_layout="NCHW",
kernel_layout="OIHW",
out_layout="NCHW",
out_dtype="float32",
)
lv2: R.Tensor((1, 6, 1, 1), dtype="float32") = R.reshape(w2, [1, 6, 1, 1])
lv3: R.Tensor((1, 6, 4, 4), dtype="float32") = R.add(lv1, lv2)
gv: R.Tensor((1, 6, 4, 4), dtype="float32") = lv3
R.output(gv)
return gv
class Conv2D2(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(3, 6, 7, bias=False)
def forward(self, input):
return self.conv(input)
@tvm.script.ir_module
class expected2:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((6, 3, 7, 7), dtype="float32"),
) -> R.Tensor((1, 6, 4, 4), dtype="float32"):
# block 0
with R.dataflow():
lv1: R.Tensor((1, 6, 4, 4), dtype="float32") = R.nn.conv2d(
input_1,
w1,
strides=[1, 1],
padding=[0, 0, 0, 0],
dilation=[1, 1],
data_layout="NCHW",
kernel_layout="OIHW",
out_layout="NCHW",
out_dtype="float32",
)
gv: R.Tensor((1, 6, 4, 4), dtype="float32") = lv1
R.output(gv)
return gv
input_info = [([1, 3, 10, 10], "float32")]
model = Conv2D1()
binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
verify_model(model, input_info, binding, expected1)
model = Conv2D1Func()
binding = {"w1": model.weight.numpy(), "w2": model.bias.numpy()}
verify_model(model, input_info, binding, expected1)
model = Conv2D2()
binding = {"w1": model.conv.weight.detach().numpy()}
verify_model(model, input_info, binding, expected2)
def test_conv2d_transpose():
class ConvTranspose2d1(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.ConvTranspose2d(3, 3, 7, bias=True)
def forward(self, input):
return self.conv(input)
class ConvTranspose2d1Func(Module):
def __init__(self):
super().__init__()
self.weight = torch.randn(size=[3, 3, 7, 7])
self.bias = torch.randn(size=[3])
def forward(self, input):
return torch.nn.functional.conv_transpose2d(input, self.weight, self.bias)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((3, 3, 7, 7), dtype="float32"),
w2: R.Tensor((3,), dtype="float32"),
) -> R.Tensor((1, 3, 16, 16), dtype="float32"):
# block 0
with R.dataflow():
lv1: R.Tensor((1, 3, 16, 16), dtype="float32") = R.nn.conv2d_transpose(
input_1,
w1,
strides=[1, 1],
padding=[0, 0, 0, 0],
output_padding=[0, 0],
dilation=[1, 1],
data_layout="NCHW",
kernel_layout="IOHW",
out_layout="NCHW",
out_dtype="float32",
)
lv2: R.Tensor((1, 3, 1, 1), dtype="float32") = R.reshape(w2, [1, 3, 1, 1])
lv3: R.Tensor((1, 3, 16, 16), dtype="float32") = R.add(lv1, lv2)
gv: R.Tensor((1, 3, 16, 16), dtype="float32") = lv3
R.output(gv)
return gv
class ConvTranspose2d2(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.ConvTranspose2d(3, 3, 7, bias=False)
def forward(self, input):
return self.conv(input)
@tvm.script.ir_module
class expected2:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((3, 3, 7, 7), dtype="float32"),
) -> R.Tensor((1, 3, 16, 16), dtype="float32"):
# block 0
with R.dataflow():
lv1: R.Tensor((1, 3, 16, 16), dtype="float32") = R.nn.conv2d_transpose(
input_1,
w1,
strides=[1, 1],
padding=[0, 0, 0, 0],
output_padding=[0, 0],
dilation=[1, 1],
data_layout="NCHW",
kernel_layout="IOHW",
out_layout="NCHW",
out_dtype="float32",
)
gv: R.Tensor((1, 3, 16, 16), dtype="float32") = lv1
R.output(gv)
return gv
input_info = [([1, 3, 10, 10], "float32")]
model = ConvTranspose2d1()
binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
verify_model(model, input_info, binding, expected1)
model = ConvTranspose2d1Func()
binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
verify_model(model, input_info, binding, expected1)
model = ConvTranspose2d2()
binding = {"w1": model.conv.weight.detach().numpy()}
verify_model(model, input_info, binding, expected2)
def test_conv3d():
class Conv3D1(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv3d(3, 6, 7, bias=True)
def forward(self, input):
return self.conv(input)
class Conv3D1Func(Module):
def __init__(self):
super().__init__()
self.weight = torch.randn(size=[6, 3, 7, 7, 7])
self.bias = torch.randn(size=[6])
def forward(self, input):
return torch.nn.functional.conv3d(input, self.weight, self.bias)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32"),
w1: R.Tensor((6, 3, 7, 7, 7), dtype="float32"),
w2: R.Tensor((6,), dtype="float32"),
) -> R.Tensor((1, 6, 4, 4, 4), dtype="float32"):
# block 0
with R.dataflow():
lv1: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = R.nn.conv3d(
input_1,
w1,
strides=[1],
padding=[0, 0, 0],
dilation=[1],
data_layout="NCDHW",
kernel_layout="OIDHW",
out_layout="NCDHW",
out_dtype="float32",
)
lv2: R.Tensor((1, 6, 1, 1, 1), dtype="float32") = R.reshape(w2, [1, 6, 1, 1, 1])
lv3: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = R.add(lv1, lv2)
gv: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = lv3
R.output(gv)
return gv
class Conv3D2(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv3d(3, 6, 7, bias=False)
def forward(self, input):
return self.conv(input)
@tvm.script.ir_module
class expected2:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32"),
w1: R.Tensor((6, 3, 7, 7, 7), dtype="float32"),
) -> R.Tensor((1, 6, 4, 4, 4), dtype="float32"):
# block 0
with R.dataflow():
lv1: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = R.nn.conv3d(
input_1,
w1,
strides=[1],
padding=[0, 0, 0],
dilation=[1],
data_layout="NCDHW",
kernel_layout="OIDHW",
out_layout="NCDHW",
out_dtype="float32",
)
gv: R.Tensor((1, 6, 4, 4, 4), dtype="float32") = lv1
R.output(gv)
return gv
input_info = [([1, 3, 10, 10, 10], "float32")]
model = Conv3D1()
binding = {"w1": model.conv.weight.detach().numpy(), "w2": model.conv.bias.detach().numpy()}
verify_model(model, input_info, binding, expected1)
model = Conv3D1Func()
binding = {"w1": model.weight.detach().numpy(), "w2": model.bias.detach().numpy()}
verify_model(model, input_info, binding, expected1)
model = Conv3D2()
binding = {"w1": model.conv.weight.detach().numpy()}
verify_model(model, input_info, binding, expected2)
def test_pad():
class PadModel(torch.nn.Module):
def __init__(self, pad, mode="constant", value=0.0):
super().__init__()
self.pad = pad
self.mode = mode
self.value = value
def forward(self, x):
if self.mode == "constant":
return torch.nn.functional.pad(x, self.pad, mode=self.mode, value=self.value)
else:
return torch.nn.functional.pad(x, self.pad, mode=self.mode)
@tvm.script.ir_module
class expected_constant:
@R.function
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 14, 12), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 14, 12), dtype="float32") = R.nn.pad(
x,
pad_width=[0, 0, 0, 0, 2, 2, 1, 1],
pad_mode="constant",
pad_value=0.0,
)
gv: R.Tensor((1, 3, 14, 12), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class expected_reflect:
@R.function
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 14, 12), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 14, 12), dtype="float32") = R.nn.pad(
x,
pad_width=[0, 0, 0, 0, 2, 2, 1, 1],
pad_mode="reflect",
pad_value=0.0,
)
gv: R.Tensor((1, 3, 14, 12), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class expected_replicate:
@R.function
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 14, 12), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 14, 12), dtype="float32") = R.nn.pad(
x,
pad_width=[0, 0, 0, 0, 2, 2, 1, 1],
pad_mode="replicate",
pad_value=0.0,
)
gv: R.Tensor((1, 3, 14, 12), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class expected_circular:
@R.function
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 14, 12), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 14, 12), dtype="float32") = R.nn.pad(
x,
pad_width=[0, 0, 0, 0, 2, 2, 1, 1],
pad_mode="circular",
pad_value=0.0,
)
gv: R.Tensor((1, 3, 14, 12), dtype="float32") = lv
R.output(gv)
return gv
input_infos = [([1, 3, 10, 10], "float32")]
verify_model(PadModel(pad=[1, 1, 2, 2]), input_infos, {}, expected_constant)
verify_model(PadModel(pad=[1, 1, 2, 2], mode="reflect"), input_infos, {}, expected_reflect)
verify_model(PadModel(pad=[1, 1, 2, 2], mode="replicate"), input_infos, {}, expected_replicate)
verify_model(PadModel(pad=[1, 1, 2, 2], mode="circular"), input_infos, {}, expected_circular)
def test_pixel_shuffle():
class PixelShuffle1(torch.nn.Module):
def __init__(self, upscale_factor=2):
super().__init__()
self.pixel_shuffle = torch.nn.PixelShuffle(upscale_factor)
def forward(self, x):
return self.pixel_shuffle(x)
class PixelShuffle2(torch.nn.Module):
def __init__(self, upscale_factor=2):
super().__init__()
self.upscale_factor = upscale_factor
def forward(self, x):
return torch.nn.functional.pixel_shuffle(x, self.upscale_factor)
@tvm.script.ir_module
class expected:
@R.function
def main(inp_0: R.Tensor((1, 8, 10, 15), dtype="float32")) -> R.Tensor(
(1, 2, 20, 30), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 2, 20, 30), dtype="float32") = R.nn.pixel_shuffle(
inp_0, upscale_factor=2
)
gv: R.Tensor((1, 2, 20, 30), dtype="float32") = lv
R.output(gv)
return gv
input_infos = [([1, 8, 10, 15], "float32")]
verify_model(PixelShuffle1(2), input_infos, {}, expected)
verify_model(PixelShuffle2(2), input_infos, {}, expected)
def test_linear():
# nn.Linear
class Dense1(Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(10, 7, bias=True)
def forward(self, input):
return self.linear(input)
class Dense1Func(Module):
def __init__(self):
super().__init__()
self.weight = torch.randn(size=[7, 10])
self.bias = torch.randn(size=[7])
def forward(self, input):
return torch.nn.functional.linear(input, self.weight, self.bias)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((7, 10), dtype="float32"),
w2: R.Tensor((7,), dtype="float32"),
) -> R.Tensor((1, 3, 10, 7), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((10, 7), dtype="float32") = R.permute_dims(w1, axes=None)
lv1: R.Tensor((1, 3, 10, 7), dtype="float32") = R.matmul(
input_1, lv, out_dtype="float32"
)
lv2: R.Tensor((1, 3, 10, 7), dtype="float32") = R.add(lv1, w2)
gv: R.Tensor((1, 3, 10, 7), dtype="float32") = lv2
R.output(gv)
return gv
class Dense2(Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(10, 7, bias=False)
def forward(self, input):
return self.linear(input)
@tvm.script.ir_module
class expected2:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((7, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 7), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((10, 7), dtype="float32") = R.permute_dims(w1, axes=None)
lv1: R.Tensor((1, 3, 10, 7), dtype="float32") = R.matmul(
input_1, lv, out_dtype="float32"
)
gv: R.Tensor((1, 3, 10, 7), dtype="float32") = lv1
R.output(gv)
return gv
input_info = [([1, 3, 10, 10], "float32")]
model = Dense1()
binding = {"w1": model.linear.weight.detach().numpy(), "w2": model.linear.bias.detach().numpy()}
verify_model(model, input_info, binding, expected1)
model = Dense1Func()
binding = {"w1": model.weight.numpy(), "w2": model.bias.numpy()}
verify_model(model, input_info, binding, expected1)
model = Dense2()
binding = {"w1": model.linear.weight.detach().numpy()}
verify_model(model, input_info, binding, expected2)
# matmul
class MatMul1(Module):
def __init__(self):
super().__init__()
def forward(self, x, y):
return torch.matmul(x, y)
@tvm.script.ir_module
class expected3:
@R.function
def main(
input_1: R.Tensor((10, 10), dtype="float32"),
input_2: R.Tensor((10, 10), dtype="float32"),
) -> R.Tensor((10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((10, 10), dtype="float32") = R.matmul(
input_1, input_2, out_dtype="float32"
)
gv: R.Tensor((10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(
MatMul1(),
[([10, 10], "float32"), ([10, 10], "float32")],
{},
expected3,
)
def test_bmm():
class BMM(Module):
def __init__(self):
super().__init__()
def forward(self, x, y):
return torch.bmm(x, y)
@tvm.script.ir_module
class Expected:
@R.function
def main(
input_1: R.Tensor((4, 128, 256), dtype="float32"),
input_2: R.Tensor((4, 256, 512), dtype="float32"),
) -> R.Tensor((4, 128, 512), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((4, 128, 512), dtype="float32") = R.matmul(
input_1, input_2, out_dtype="float32"
)
gv: R.Tensor((4, 128, 512), dtype="float32") = lv
R.output(gv)
return gv
verify_model(
BMM(),
[((4, 128, 256), "float32"), ((4, 256, 512), "float32")],
{},
Expected,
)
def test_baddbmm():
class BAddBMM1(Module):
def __init__(self):
super().__init__()
def forward(self, c, x, y):
return torch.baddbmm(c, x, y)
class BAddBMM2(Module):
def __init__(self):
super().__init__()
def forward(self, c, x, y):
return torch.baddbmm(c, x, y, alpha=2, beta=0)
@tvm.script.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((4, 128, 512), dtype="float32"),
inp_1: R.Tensor((4, 128, 256), dtype="float32"),
inp_2: R.Tensor((4, 256, 512), dtype="float32"),
) -> R.Tensor((4, 128, 512), dtype="float32"):
with R.dataflow():
lv: R.Tensor((4, 128, 512), dtype="float32") = R.matmul(inp_1, inp_2)
lv1: R.Tensor((4, 128, 512), dtype="float32") = R.add(lv, inp_0)
gv: R.Tensor((4, 128, 512), dtype="float32") = lv1
R.output(gv)
return gv
@tvm.script.ir_module
class Expected2:
@R.function
def main(
inp_0: R.Tensor((4, 128, 512), dtype="float32"),
inp_1: R.Tensor((4, 128, 256), dtype="float32"),
inp_2: R.Tensor((4, 256, 512), dtype="float32"),
) -> R.Tensor((4, 128, 512), dtype="float32"):
with R.dataflow():
lv: R.Tensor((4, 128, 512), dtype="float32") = R.matmul(inp_1, inp_2)
lv1: R.Tensor((4, 128, 512), dtype="float32") = R.multiply(
lv, R.const(2, "float32")
)
gv: R.Tensor((4, 128, 512), dtype="float32") = lv1
R.output(gv)
return gv
verify_model(
BAddBMM1(),
[((4, 128, 512), "float32"), ((4, 128, 256), "float32"), ((4, 256, 512), "float32")],
{},
Expected1,
)
verify_model(
BAddBMM2(),
[((4, 128, 512), "float32"), ((4, 128, 256), "float32"), ((4, 256, 512), "float32")],
{},
Expected2,
)
def test_einsum():
class Einsum1(Module):
def __init__(self):
super().__init__()
def forward(self, x):
return torch.einsum("ii", x)
class Einsum2(Module):
def __init__(self):
super().__init__()
def forward(self, x, y):
return torch.einsum("i,j->ij", x, y)
@tvm.script.ir_module
class Expected1:
@R.function
def main(inp_0: R.Tensor((4, 4), dtype="float32")) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((), dtype="float32") = R.einsum((inp_0,), subscripts="ii")
gv: R.Tensor((), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class Expected2:
@R.function
def main(
inp_0: R.Tensor((5,), dtype="float32"), inp_1: R.Tensor((4,), dtype="float32")
) -> R.Tensor((5, 4), dtype="float32"):
with R.dataflow():
lv: R.Tensor((5, 4), dtype="float32") = R.einsum(
(inp_0, inp_1), subscripts="i,j->ij"
)
gv: R.Tensor((5, 4), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Einsum1(), [([4, 4], "float32")], {}, Expected1)
verify_model(Einsum2(), [([5], "float32"), ([4], "float32")], {}, Expected2)
def test_outer():
class Outer(torch.nn.Module):
def forward(self, x, y):
return torch.outer(x, y)
@tvm.script.ir_module
class expected:
@R.function
def main(
a: R.Tensor((3,), dtype="float32"), b: R.Tensor((4,), dtype="float32")
) -> R.Tensor((3, 4), dtype="float32"):
with R.dataflow():
lv: R.Tensor((3, 4), dtype="float32") = R.outer(a, b)
gv: R.Tensor((3, 4), dtype="float32") = lv
R.output(gv)
return gv
input_infos = [([3], "float32"), ([4], "float32")]
verify_model(Outer(), input_infos, {}, expected)
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
def test_softplus():
import torch
from torch.nn import Module
torch.set_grad_enabled(False)
class Softplus0(torch.nn.Module):
def __init__(self):
super().__init__()
self.softplus = torch.nn.Softplus(1.0, 20.0)
def forward(self, x):
return self.softplus(x)
class Softplus1(Module):
def forward(self, input):
return torch.nn.functional.softplus(input, 1.0, 20.0)
@tvm.script.ir_module
class expected:
@R.function
def main(inp_0: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((10, 10), dtype="float32") = R.nn.softplus(
inp_0, beta=1.0, threshold=20.0
)
gv: R.Tensor((10, 10), dtype="float32") = lv
R.output(gv)
return gv
input_info = [([10, 10], "float32")]
verify_model(Softplus0(), input_info, {}, expected)
verify_model(Softplus1(), input_info, {}, expected)
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
def test_leakyrelu():
import torch
from torch.nn import Module
torch.set_grad_enabled(False)
class LeakyReLU0(Module):
def __init__(self):
super().__init__()
self.leakyrelu = torch.nn.LeakyReLU(0.02)
def forward(self, input):
return self.leakyrelu(input)
class LeakyReLU1(Module):
def forward(self, input):
return torch.nn.functional.leaky_relu(input, 0.02)
@tvm.script.ir_module
class expected:
@R.function
def main(input_1: R.Tensor((10, 10), dtype="float32")) -> R.Tensor(
(10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((10, 10), dtype="float32") = R.nn.leakyrelu(input_1, 0.02)
gv: R.Tensor((10, 10), dtype="float32") = lv
R.output(gv)
return gv
input_info = [([10, 10], "float32")]
verify_model(LeakyReLU0(), input_info, {}, expected)
verify_model(LeakyReLU1(), input_info, {}, expected)
def test_prelu():
class Prelu1(Module):
def __init__(self, num_parameters=1, alpha=0.25):
super().__init__()
self.prelu = torch.nn.PReLU(num_parameters=num_parameters, init=alpha)
def forward(self, x):
return self.prelu(x)
class Prelu2(torch.nn.Module):
def __init__(self):
super().__init__()
self.alpha = torch.nn.Parameter(torch.tensor([0.25]))
def forward(self, x):
return torch.nn.functional.prelu(x, self.alpha)
@tvm.script.ir_module
class expected:
@R.function
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.prelu(
x, R.const([0.25], dtype="float32"), axis=1
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
input_info = [([1, 3, 10, 10], "float32")]
verify_model(Prelu1(), input_info, {}, expected)
verify_model(Prelu2(), input_info, {}, expected)
def test_maxpool1d():
input_info = [([1, 3, 10], "float32")]
class MaxPool1d(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.MaxPool1d(kernel_size=2)
def forward(self, input):
return self.pool(input)
class MaxPool1d_functional(Module):
def __init__(self):
super().__init__()
def forward(self, input):
return torch.nn.functional.max_pool1d(input, kernel_size=2)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tensor(
(1, 3, 5), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 5), dtype="float32") = R.nn.max_pool1d(
input_1,
pool_size=[2],
strides=[2],
dilation=[1],
padding=[0, 0],
layout="NCW",
out_layout="NCW",
)
gv: R.Tensor((1, 3, 5), dtype="float32") = lv
R.output(gv)
return gv
class MaxPool1d2(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.MaxPool1d(kernel_size=3, stride=1, padding=1)
def forward(self, input):
return self.pool(input)
@tvm.script.ir_module
class expected2:
@R.function
def main(input_1: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10), dtype="float32") = R.nn.max_pool1d(
input_1,
pool_size=[3],
strides=[1],
dilation=[1],
padding=[1, 1],
layout="NCW",
out_layout="NCW",
)
gv: R.Tensor((1, 3, 10), dtype="float32") = lv
R.output(gv)
return gv
class MaxPool1d3(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.MaxPool1d(kernel_size=3, stride=2, dilation=2)
def forward(self, input):
return self.pool(input)
@tvm.script.ir_module
class expected3:
@R.function
def main(input_1: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tensor(
(1, 3, 3), dtype="float32"
): # Corrected here
with R.dataflow():
lv: R.Tensor((1, 3, 3), dtype="float32") = R.nn.max_pool1d(
input_1,
pool_size=[3],
strides=[2],
dilation=[2],
padding=[0, 0],
layout="NCW",
out_layout="NCW",
)
gv: R.Tensor((1, 3, 3), dtype="float32") = lv
R.output(gv)
return gv
verify_model(MaxPool1d(), input_info, {}, expected1)
verify_model(MaxPool1d_functional(), input_info, {}, expected1)
verify_model(MaxPool1d2(), input_info, {}, expected2)
verify_model(MaxPool1d3(), input_info, {}, expected3)
def test_maxpool2d():
input_info = [([1, 3, 10, 10], "float32")]
class MaxPool2d(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.MaxPool2d(kernel_size=[1, 1])
def forward(self, input):
return self.pool(input)
class MaxPool2d_functional(Module):
def __init__(self):
super().__init__()
def forward(self, input):
return torch.nn.functional.max_pool2d(input, kernel_size=[1, 1])
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.max_pool2d(
input_1,
pool_size=[1, 1],
strides=[1, 1],
dilation=[1, 1],
padding=[0, 0, 0, 0],
layout="NCHW",
out_layout="NCHW",
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
class MaxPool2d2(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.MaxPool2d(kernel_size=[2, 2], dilation=[2, 3])
def forward(self, input):
return self.pool(input)
@tvm.script.ir_module
class expected2:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 4, 4), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 4, 4), dtype="float32") = R.nn.max_pool2d(
input_1,
pool_size=[2, 2],
strides=[2, 2],
dilation=[2, 3],
padding=[0, 0, 0, 0],
layout="NCHW",
out_layout="NCHW",
)
gv: R.Tensor((1, 3, 4, 4), dtype="float32") = lv
R.output(gv)
return gv
class MaxPool2d3(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.MaxPool2d(kernel_size=[4, 4], padding=2, stride=2)
def forward(self, input):
return self.pool(input)
@tvm.script.ir_module
class expected3:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 6, 6), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 6, 6), dtype="float32") = R.nn.max_pool2d(
input_1,
pool_size=[4, 4],
strides=[2, 2],
dilation=[1, 1],
padding=[2, 2, 2, 2],
layout="NCHW",
out_layout="NCHW",
)
gv: R.Tensor((1, 3, 6, 6), dtype="float32") = lv
R.output(gv)
return gv
verify_model(MaxPool2d(), input_info, {}, expected1)
verify_model(MaxPool2d_functional(), input_info, {}, expected1)
verify_model(MaxPool2d2(), input_info, {}, expected2)
verify_model(MaxPool2d3(), input_info, {}, expected3)
def test_maxpool3d():
input_info = [([1, 3, 10, 10, 10], "float32")]
class MaxPool3d(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.MaxPool3d(kernel_size=[1, 1, 1])
def forward(self, input):
return self.pool(input)
class MaxPool3d_functional(Module):
def __init__(self):
super().__init__()
def forward(self, input):
return torch.nn.functional.max_pool3d(input, kernel_size=[1, 1, 1])
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10, 10), dtype="float32") = R.nn.max_pool3d(
input_1,
pool_size=[1, 1, 1],
strides=[1, 1, 1],
dilation=[1, 1, 1],
padding=[0, 0, 0, 0, 0, 0],
layout="NCDHW",
out_layout="NCDHW",
)
gv: R.Tensor((1, 3, 10, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
class MaxPool3d2(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.MaxPool3d(kernel_size=[2, 2, 2], dilation=[1, 2, 2])
def forward(self, input):
return self.pool(input)
@tvm.script.ir_module
class expected2:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 5, 4, 4), dtype="float32"
): # Fixed here
with R.dataflow():
lv: R.Tensor((1, 3, 5, 4, 4), dtype="float32") = R.nn.max_pool3d(
input_1,
pool_size=[2, 2, 2],
strides=[2, 2, 2],
dilation=[1, 2, 2],
padding=[0, 0, 0, 0, 0, 0],
layout="NCDHW",
out_layout="NCDHW",
)
gv: R.Tensor((1, 3, 5, 4, 4), dtype="float32") = lv
R.output(gv)
return gv
class MaxPool3d3(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.MaxPool3d(kernel_size=[3, 3, 3], padding=1, stride=2)
def forward(self, input):
return self.pool(input)
@tvm.script.ir_module
class expected3:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 5, 5, 5), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 5, 5, 5), dtype="float32") = R.nn.max_pool3d(
input_1,
pool_size=[3, 3, 3],
strides=[2, 2, 2],
dilation=[1, 1, 1],
padding=[1, 1, 1, 1, 1, 1],
layout="NCDHW",
out_layout="NCDHW",
)
gv: R.Tensor((1, 3, 5, 5, 5), dtype="float32") = lv
R.output(gv)
return gv
verify_model(MaxPool3d(), input_info, {}, expected1)
verify_model(MaxPool3d_functional(), input_info, {}, expected1)
verify_model(MaxPool3d2(), input_info, {}, expected2)
verify_model(MaxPool3d3(), input_info, {}, expected3)
def test_avgpool1d():
input_info = [([1, 3, 10], "float32")]
class AvgPool1d(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.AvgPool1d(kernel_size=1)
def forward(self, input):
return self.pool(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10), dtype="float32"
):
with R.dataflow():
lv = R.nn.avg_pool1d(
input_1,
pool_size=[1],
strides=[1],
dilation=[1],
padding=[0, 0],
ceil_mode=False,
count_include_pad=True,
layout="NCW",
out_layout="NCW",
)
gv = lv
R.output(gv)
return gv
class AvgPool1d2(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.AvgPool1d(kernel_size=4, stride=2, padding=2, ceil_mode=True)
def forward(self, input):
return self.pool(input)
class AvgPool1d3(Module):
def forward(self, input):
return torch.nn.functional.avg_pool1d(
input, kernel_size=4, stride=2, padding=2, ceil_mode=True
)
@tvm.script.ir_module
class expected2:
@R.function
def main(input_1: R.Tensor((1, 3, 10), dtype="float32")):
with R.dataflow():
lv = R.nn.avg_pool1d(
input_1,
pool_size=[4],
strides=[2],
dilation=[1],
padding=[2, 2],
ceil_mode=True,
count_include_pad=True,
layout="NCW",
out_layout="NCW",
)
gv = lv
R.output(gv)
return gv
class AvgPool1d4(Module):
def forward(self, input):
return torch.nn.functional.avg_pool1d(input, kernel_size=2)
@tvm.script.ir_module
class expected3:
@R.function
def main(input_1: R.Tensor((1, 3, 10), dtype="float32")):
with R.dataflow():
lv = R.nn.avg_pool1d(
input_1,
pool_size=[2],
strides=[2],
dilation=[1],
padding=[0, 0],
ceil_mode=False,
count_include_pad=True,
layout="NCW",
out_layout="NCW",
)
gv = lv
R.output(gv)
return gv
verify_model(AvgPool1d(), input_info, {}, expected1)
verify_model(AvgPool1d2(), input_info, {}, expected2)
verify_model(AvgPool1d3(), input_info, {}, expected2)
verify_model(AvgPool1d4(), input_info, {}, expected3)
def test_avgpool2d():
input_info = [([1, 3, 10, 10], "float32")]
class AvgPool2d(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.AvgPool2d(kernel_size=[1, 1])
def forward(self, input):
return self.pool(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.avg_pool2d(
input_1,
pool_size=[1, 1],
strides=[1, 1],
dilation=[1, 1],
padding=[0, 0, 0, 0],
count_include_pad=True,
layout="NCHW",
out_layout="NCHW",
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
class AvgPool2d2(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.AvgPool2d(kernel_size=[4, 4], stride=2, padding=2, ceil_mode=True)
def forward(self, input):
return self.pool(input)
class AvgPool2d3(Module):
def forward(self, input):
return torch.nn.functional.avg_pool2d(
input, kernel_size=[4, 4], stride=2, padding=2, ceil_mode=True
)
@tvm.script.ir_module
class expected2:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")):
with R.dataflow():
lv = R.nn.avg_pool2d(
input_1,
pool_size=[4, 4],
strides=[2, 2],
dilation=[1, 1],
padding=[2, 2, 2, 2],
ceil_mode=True,
count_include_pad=True,
layout="NCHW",
out_layout="NCHW",
)
gv = lv
R.output(gv)
return gv
class AvgPool2d4(Module):
def forward(self, input):
return torch.nn.functional.avg_pool2d(input, kernel_size=[2, 1], divisor_override=2)
@tvm.script.ir_module
class expected3:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")):
with R.dataflow():
lv = R.nn.avg_pool2d(
input_1,
pool_size=[2, 1],
strides=[2, 1],
dilation=[1, 1],
padding=[0, 0, 0, 0],
ceil_mode=False,
count_include_pad=True,
layout="NCHW",
out_layout="NCHW",
)
gv = lv
R.output(gv)
return gv
verify_model(AvgPool2d(), input_info, {}, expected1)
verify_model(AvgPool2d2(), input_info, {}, expected2)
verify_model(AvgPool2d3(), input_info, {}, expected2)
verify_model(AvgPool2d4(), input_info, {}, expected3)
def test_avgpool3d():
input_info = [([1, 3, 8, 8, 8], "float32")]
class AvgPool3d(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.AvgPool3d(kernel_size=[1, 1, 1])
def forward(self, input):
return self.pool(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")) -> R.Tensor(
(1, 3, 8, 8, 8), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 8, 8, 8), dtype="float32") = R.nn.avg_pool3d(
input_1,
pool_size=[1, 1, 1],
strides=[1, 1, 1],
dilation=[1, 1, 1],
padding=[0, 0, 0, 0, 0, 0],
ceil_mode=False,
count_include_pad=True,
layout="NCDHW",
out_layout="NCDHW",
)
gv: R.Tensor((1, 3, 8, 8, 8), dtype="float32") = lv
R.output(gv)
return gv
class AvgPool3d2(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.AvgPool3d(
kernel_size=[3, 3, 3], stride=2, padding=1, ceil_mode=True
)
def forward(self, input):
return self.pool(input)
class AvgPool3d3(Module):
def forward(self, input):
return torch.nn.functional.avg_pool3d(
input, kernel_size=[3, 3, 3], stride=2, padding=1, ceil_mode=True
)
@tvm.script.ir_module
class expected2:
@R.function
def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")):
with R.dataflow():
lv = R.nn.avg_pool3d(
input_1,
pool_size=[3, 3, 3],
strides=[2, 2, 2],
dilation=[1, 1, 1],
padding=[1, 1, 1, 1, 1, 1],
ceil_mode=True,
count_include_pad=True,
layout="NCDHW",
out_layout="NCDHW",
)
gv = lv
R.output(gv)
return gv
class AvgPool3d4(Module):
def forward(self, input):
return torch.nn.functional.avg_pool3d(input, kernel_size=[2, 1, 2], stride=[2, 1, 2])
@tvm.script.ir_module
class expected3:
@R.function
def main(input_1: R.Tensor((1, 3, 8, 8, 8), dtype="float32")):
with R.dataflow():
lv = R.nn.avg_pool3d(
input_1,
pool_size=[2, 1, 2],
strides=[2, 1, 2],
dilation=[1, 1, 1],
padding=[0, 0, 0, 0, 0, 0],
ceil_mode=False,
count_include_pad=True,
layout="NCDHW",
out_layout="NCDHW",
)
gv = lv
R.output(gv)
return gv
verify_model(AvgPool3d(), input_info, {}, expected1)
verify_model(AvgPool3d2(), input_info, {}, expected2)
verify_model(AvgPool3d3(), input_info, {}, expected2)
verify_model(AvgPool3d4(), input_info, {}, expected3)
def test_adaptive_avgpool1d():
input_info = [([1, 3, 16], "float32")]
class AdaptiveAvgPool1d0(torch.nn.Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.AdaptiveAvgPool1d(8)
def forward(self, input):
return self.pool(input)
class AdaptiveAvgPool1d1(torch.nn.Module):
def forward(self, input):
return torch.nn.functional.adaptive_avg_pool1d(input, 8)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 16), dtype="float32")) -> R.Tensor(
(1, 3, 8), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 8), dtype="float32") = R.nn.adaptive_avg_pool1d(
input_1, output_size=[8], layout="NCW", out_layout="NCW"
)
gv: R.Tensor((1, 3, 8), dtype="float32") = lv
R.output(gv)
return gv
verify_model(AdaptiveAvgPool1d0(), input_info, {}, expected1)
verify_model(AdaptiveAvgPool1d1(), input_info, {}, expected1)
def test_adaptive_avgpool2d():
input_info = [([1, 3, 10, 10], "float32")]
class AdaptiveAvgPool2d0(Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.AdaptiveAvgPool2d([10, 10])
def forward(self, input):
return self.pool(input)
class AdaptiveAvgPool2d1(Module):
def forward(self, input):
return torch.nn.functional.adaptive_avg_pool2d(input, [10, 10])
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.adaptive_avg_pool2d(
input_1, output_size=[10, 10], layout="NCHW", out_layout="NCHW"
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(AdaptiveAvgPool2d0(), input_info, {}, expected1)
verify_model(AdaptiveAvgPool2d1(), input_info, {}, expected1)
def test_adaptive_avgpool3d():
input_info = [([1, 3, 16, 16, 16], "float32")]
class AdaptiveAvgPool3d0(torch.nn.Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.AdaptiveAvgPool3d((8, 8, 8))
def forward(self, input):
return self.pool(input)
class AdaptiveAvgPool3d1(torch.nn.Module):
def forward(self, input):
return torch.nn.functional.adaptive_avg_pool3d(input, (8, 8, 8))
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 16, 16, 16), dtype="float32")) -> R.Tensor(
(1, 3, 8, 8, 8), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 8, 8, 8), dtype="float32") = R.nn.adaptive_avg_pool3d(
input_1, output_size=[8, 8, 8], layout="NCDHW", out_layout="NCDHW"
)
gv: R.Tensor((1, 3, 8, 8, 8), dtype="float32") = lv
R.output(gv)
return gv
verify_model(AdaptiveAvgPool3d0(), input_info, {}, expected1)
verify_model(AdaptiveAvgPool3d1(), input_info, {}, expected1)
def test_flatten():
input_info = [([1, 3, 10, 10], "float32")]
class Flatten(Module):
def __init__(self):
super().__init__()
self.f = torch.nn.Flatten(2, -1)
def forward(self, input):
return self.f(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 100), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 100), dtype="float32") = R.reshape(input_1, (1, 3, 100))
gv: R.Tensor((1, 3, 100), dtype="float32") = lv
R.output(gv)
return gv
# call_module
verify_model(Flatten(), input_info, {}, expected1)
# call_method
verify_model(torch.nn.Flatten(2, -1), input_info, {}, expected1)
def test_batchnorm2d():
input_info = [([1, 3, 10, 10], "float32")]
class BatchNorm2d(Module):
def __init__(self):
super().__init__()
self.bn = torch.nn.BatchNorm2d(3)
def forward(self, input):
return self.bn(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((3,), dtype="float32"),
w2: R.Tensor((3,), dtype="float32"),
w3: R.Tensor((3,), dtype="float32"),
w4: R.Tensor((3,), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tuple(
R.Tensor((1, 3, 10, 10), dtype="float32"),
R.Tensor((3,), dtype="float32"),
R.Tensor((3,), dtype="float32"),
) = R.nn.batch_norm(
input_1,
w1,
w2,
w3,
w4,
axis=1,
epsilon=1e-05,
center=True,
scale=True,
)
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[0]
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv1
R.output(gv)
return gv
model = BatchNorm2d()
binding = {
"w1": model.bn.weight.detach().numpy(),
"w2": model.bn.bias.detach().numpy(),
"w3": model.bn.running_mean.detach().numpy(),
"w4": model.bn.running_var.detach().numpy(),
}
verify_model(BatchNorm2d(), input_info, binding, expected1)
def test_embedding():
input_info = [([4], "int64")]
class Embedding(Module):
def __init__(self):
super().__init__()
self.embedding = torch.nn.Embedding(10, 3)
def forward(self, input):
return self.embedding(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((4,), dtype="int64"), w1: R.Tensor((10, 3), dtype="float32")
) -> R.Tensor((4, 3), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((4,), dtype="int32") = R.astype(input_1, dtype="int32")
lv1: R.Tensor((4, 3), dtype="float32") = R.take(w1, lv, axis=0)
gv: R.Tensor((4, 3), dtype="float32") = lv1
R.output(gv)
return gv
model = Embedding()
binding = {"w1": model.embedding.weight.detach().numpy()}
verify_model(model, input_info, binding, expected1)
def test_stochastic_depth():
torchvision = pytest.importorskip("torchvision")
input_info = [([1, 3, 10, 10], "float32")]
class StochasticDepth1(Module):
def __init__(self):
super().__init__()
self.stochastic_depth = torchvision.ops.StochasticDepth(0.5, mode="row")
def forward(self, x):
return self.stochastic_depth(x)
class StochasticDepth2(Module):
def forward(self, x):
return torchvision.ops.stochastic_depth(x, 0.5, mode="row", training=False)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = input_1
R.output(gv)
return gv
verify_model(StochasticDepth1(), input_info, {}, expected1)
verify_model(StochasticDepth2(), input_info, {}, expected1)
def test_layernorm():
input_info = [([1, 3, 10, 10], "float32")]
class LayerNorm(Module):
def __init__(self):
super().__init__()
self.ln = torch.nn.LayerNorm((10, 10))
def forward(self, input):
return self.ln(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((10, 10), dtype="float32"),
w2: R.Tensor((10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm(
input_1,
w1,
w2,
axes=[-2, -1],
epsilon=1e-05,
center=True,
scale=True,
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
model = LayerNorm()
binding = {
"w1": model.ln.weight.detach().numpy(),
"w2": model.ln.bias.detach().numpy(),
}
verify_model(LayerNorm(), input_info, binding, expected1)
def test_functional_layernorm():
import numpy as np
input_info = [([1, 3, 10, 10], "float32")]
class LayerNorm(Module):
def __init__(self, shape):
super().__init__()
self.weight = torch.nn.Parameter(torch.ones(shape))
self.bias = torch.nn.Parameter(torch.zeros(shape))
def forward(self, input):
return torch.nn.functional.layer_norm(
input, self.weight.shape, self.weight, self.bias, 1e-5
)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((10, 10), dtype="float32"),
w2: R.Tensor((10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm(
input_1,
w1,
w2,
axes=[-2, -1],
epsilon=1e-05,
center=True,
scale=True,
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
model = LayerNorm((10, 10))
binding = {
"w1": model.weight.detach().numpy(),
"w2": model.bias.detach().numpy(),
}
verify_model(model, input_info, binding, expected1)
class LayerNorm2(Module):
def __init__(self, shape):
super().__init__()
self.shape = shape
self.weight = None
self.bias = None
def forward(self, input):
return torch.nn.functional.layer_norm(input, self.shape, self.weight, self.bias, 1e-5)
@tvm.script.ir_module
class expected2:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm(
input_1,
gamma=relax.const(np.ones((10, 10)), dtype="float32"),
beta=relax.const(np.zeros((10, 10)), dtype="float32"),
axes=[-2, -1],
epsilon=1e-05,
center=True,
scale=True,
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
model = LayerNorm2((10, 10))
binding = {}
verify_model(model, input_info, binding, expected2)
class LayerNorm3(Module):
def __init__(self, shape):
super().__init__()
self.shape = shape
self.weight = torch.nn.Parameter(torch.ones(shape))
self.bias = torch.nn.Parameter(torch.zeros(shape))
def forward(self, input):
return torch.nn.functional.layer_norm(input, self.shape, self.weight, self.bias, 1e-5)
@tvm.script.ir_module
class expected3:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor([10, 10], dtype="float32"),
w2: R.Tensor([10, 10], dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.layer_norm(
input_1,
w1,
w2,
axes=[-2, -1],
epsilon=1e-05,
center=True,
scale=True,
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
model = LayerNorm3([10, 10])
binding = {
"w1": model.weight.detach().numpy(),
"w2": model.bias.detach().numpy(),
}
verify_model(model, input_info, binding, expected3)
def test_cross_entropy():
input_info = [([3, 2], "float32"), ([3], "int32")]
class CrossEntropy1(Module):
def __init__(self):
super().__init__()
self.loss = torch.nn.CrossEntropyLoss()
def forward(self, logits, targets):
return self.loss(logits, targets)
@tvm.script.ir_module
class expected1:
@R.function
def main(
inp_0: R.Tensor((3, 2), dtype="float32"), inp_1: R.Tensor((3,), dtype="int32")
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((3, 2), dtype="float32") = R.nn.log_softmax(inp_0, axis=-1)
lv1: R.Tensor((), dtype="float32") = R.nn.nll_loss(
lv, inp_1, reduction="mean", ignore_index=-100
)
gv: R.Tensor((), dtype="float32") = lv1
R.output(gv)
return gv
class CrossEntropy2(Module):
def __init__(self):
super().__init__()
self.weight = torch.nn.Parameter(torch.ones((2,)))
self.loss = torch.nn.CrossEntropyLoss(weight=self.weight)
def forward(self, logits, targets):
return self.loss(logits, targets)
@tvm.script.ir_module
class expected2:
@R.function
def main(
inp_0: R.Tensor((3, 2), dtype="float32"),
inp_1: R.Tensor((3,), dtype="int32"),
w1: R.Tensor((2,), dtype="float32"),
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((3, 2), dtype="float32") = R.nn.log_softmax(inp_0, axis=-1)
lv1: R.Tensor((), dtype="float32") = R.nn.nll_loss(
lv,
inp_1,
w1,
reduction="mean",
ignore_index=-100,
)
gv: R.Tensor((), dtype="float32") = lv1
R.output(gv)
return gv
class CrossEntropy3(Module):
def __init__(self):
super().__init__()
self.loss = torch.nn.CrossEntropyLoss(ignore_index=1, reduction="sum")
def forward(self, logits, targets):
return self.loss(logits, targets)
@tvm.script.ir_module
class expected3:
@R.function
def main(
inp_0: R.Tensor((3, 2), dtype="float32"), inp_1: R.Tensor((3,), dtype="int32")
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((3, 2), dtype="float32") = R.nn.log_softmax(inp_0, axis=-1)
lv1: R.Tensor((), dtype="float32") = R.nn.nll_loss(
lv, inp_1, reduction="sum", ignore_index=1
)
gv: R.Tensor((), dtype="float32") = lv1
R.output(gv)
return gv
verify_model(CrossEntropy1(), input_info, {}, expected1)
model = CrossEntropy2()
binding = {"w1": model.loss.weight.detach().numpy()}
verify_model(model, input_info, binding, expected2)
verify_model(CrossEntropy3(), input_info, {}, expected3)
def test_functional_cross_entropy():
input_info = [([3, 10], "float32"), ([3], "int32")]
class CrossEntropy(Module):
def forward(self, logits, targets):
return torch.nn.functional.cross_entropy(logits, targets)
@tvm.script.ir_module
class expected1:
@R.function
def main(
inp_0: R.Tensor((3, 10), dtype="float32"), inp_1: R.Tensor((3,), dtype="int32")
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((3, 10), dtype="float32") = R.nn.log_softmax(inp_0, axis=-1)
lv1: R.Tensor((), dtype="float32") = R.nn.nll_loss(
lv, inp_1, reduction="mean", ignore_index=-100
)
gv: R.Tensor((), dtype="float32") = lv1
R.output(gv)
return gv
model = CrossEntropy()
verify_model(model, input_info, {}, expected1)
def test_groupnorm():
import torch
from torch.nn import Module
torch.set_grad_enabled(False)
torch.random.manual_seed(0)
input_info = [([1, 3, 10, 10], "float32")]
class GroupNorm(Module):
def __init__(self):
super().__init__()
self.gn = torch.nn.GroupNorm(3, 3)
def forward(self, input):
return self.gn(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((3,), dtype="float32"),
w2: R.Tensor((3,), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.group_norm(
input_1,
w1,
w2,
num_groups=3,
channel_axis=1,
axes=[2, 3],
epsilon=1.0000000000000001e-05,
center=True,
scale=True,
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
model = GroupNorm()
binding = {
"w1": model.gn.weight.detach().numpy(),
"w2": model.gn.bias.detach().numpy(),
}
verify_model(model, input_info, binding, expected1)
def test_instancenorm2d():
torch.set_grad_enabled(False)
torch.random.manual_seed(0)
input_info = [([1, 3, 10, 10], "float32")]
class InstanceNorm2d(Module):
def __init__(self):
super().__init__()
self.gn = torch.nn.InstanceNorm2d(3)
def forward(self, input):
return self.gn(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((3,), dtype="float32"),
w2: R.Tensor((3,), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.instance_norm(
input_1,
w1,
w2,
channel_axis=1,
axes=[2, 3],
epsilon=1e-05,
center=True,
scale=True,
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
model = InstanceNorm2d()
binding = {
"w1": torch.ones(3).detach().numpy(),
"w2": torch.zeros(3).detach().numpy(),
}
verify_model(model, input_info, binding, expected1)
operator_binary_1 = [
(operator.add, R.add),
(operator.sub, R.subtract),
(operator.mul, R.multiply),
(operator.truediv, R.divide),
(operator.floordiv, R.floor_divide),
(torch.ops.aten.fmod, R.mod),
(operator.pow, R.power),
(operator.mod, R.floor_mod),
]
@pytest.mark.parametrize("op, relax_op", operator_binary_1)
def test_binary1(op, relax_op):
input_info1 = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")]
input_info2 = [([1, 3, 10, 10], "float32")]
class Binary1(Module):
def __init__(self, op):
super().__init__()
self.op = op
def forward(self, lhs, rhs):
return self.op(lhs, rhs)
@tvm.script.ir_module
class expected_binary1:
@R.function
def main(
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(lhs, rhs)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
class Binary2(Module):
def __init__(self, op):
super().__init__()
self.op = op
def forward(self, lhs):
return self.op(lhs, 1.0)
@tvm.script.ir_module
class expected_binary2:
@R.function
def main(
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(lhs, R.const(1.0))
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Binary1(op), input_info1, {}, expected_binary1)
verify_model(Binary2(op), input_info2, {}, expected_binary2)
operator_binary_2 = [
(operator.eq, R.equal),
(operator.ne, R.not_equal),
(operator.lt, R.less),
(operator.le, R.less_equal),
(operator.gt, R.greater),
(operator.ge, R.greater_equal),
]
@pytest.mark.parametrize("op, relax_op", operator_binary_2)
def test_binary2(op, relax_op):
input_info1 = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")]
input_info2 = [([1, 3, 10, 10], "float32")]
class Binary1(Module):
def __init__(self, op):
super().__init__()
self.op = op
def forward(self, lhs, rhs):
return self.op(lhs, rhs)
@tvm.script.ir_module
class expected_binary1:
@R.function
def main(
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="bool"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(lhs, rhs)
gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv
R.output(gv)
return gv
class Binary2(Module):
def __init__(self, op):
super().__init__()
self.op = op
def forward(self, lhs):
return self.op(lhs, 1.0)
@tvm.script.ir_module
class expected_binary2:
@R.function
def main(
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="bool"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(lhs, R.const(1.0))
gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv
R.output(gv)
return gv
verify_model(Binary1(op), input_info1, {}, expected_binary1)
verify_model(Binary2(op), input_info2, {}, expected_binary2)
operator_binary_3 = [
(torch.ops.aten.bitwise_or_, R.bitwise_or),
(torch.ops.aten.bitwise_or, R.bitwise_or),
(operator.lshift, R.left_shift),
(operator.rshift, R.right_shift),
(operator.and_, R.bitwise_and),
(operator.or_, R.bitwise_or),
(operator.xor, R.bitwise_xor),
]
@pytest.mark.parametrize("op, relax_op", operator_binary_3)
def test_binary3(op, relax_op):
input_info1 = [([1, 3, 10, 10], "int32"), ([1, 3, 10, 10], "int32")]
input_info2 = [([1, 3, 10, 10], "int32")]
class Binary1(Module):
def __init__(self, op):
super().__init__()
self.op = op
def forward(self, lhs, rhs):
return self.op(lhs, rhs)
@tvm.script.ir_module
class expected_binary1:
@R.function
def main(
lhs: R.Tensor((1, 3, 10, 10), dtype="int32"),
rhs: R.Tensor((1, 3, 10, 10), dtype="int32"),
) -> R.Tensor((1, 3, 10, 10), dtype="int32"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="int32") = relax_op(lhs, rhs)
gv: R.Tensor((1, 3, 10, 10), dtype="int32") = lv
R.output(gv)
return gv
class Binary2(Module):
def __init__(self, op):
super().__init__()
self.op = op
def forward(self, lhs):
return self.op(lhs, 1)
@tvm.script.ir_module
class expected_binary2:
@R.function
def main(
lhs: R.Tensor((1, 3, 10, 10), dtype="int32"),
) -> R.Tensor((1, 3, 10, 10), dtype="int32"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="int32") = relax_op(lhs, R.const(1))
gv: R.Tensor((1, 3, 10, 10), dtype="int32") = lv
R.output(gv)
return gv
verify_model(Binary1(op), input_info1, {}, expected_binary1)
verify_model(Binary2(op), input_info2, {}, expected_binary2)
# RSub
def test_rsub():
input_info1 = [([10, 10], "float32"), ([10, 10], "float32")]
input_info2 = [([10, 10], "float32")]
class RSub1(Module):
def forward(self, x, y):
return torch.rsub(x, y)
class RSub2(Module):
def forward(self, x):
return torch.rsub(x, 5.0)
@tvm.script.ir_module
class expected_rsub1:
@R.function
def main(
x: R.Tensor((10, 10), dtype="float32"), y: R.Tensor((10, 10), dtype="float32")
) -> R.Tensor((10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((10, 10), dtype="float32") = R.subtract(y, x)
gv: R.Tensor((10, 10), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class expected_rsub2:
@R.function
def main(x: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((10, 10), dtype="float32") = R.subtract(R.const(5.0, "float32"), x)
gv: R.Tensor((10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(RSub1(), input_info1, {}, expected_rsub1)
verify_model(RSub2(), input_info2, {}, expected_rsub2)
# IsIn
def test_isin():
input_info = [([10, 10], "float32"), ([8], "float32")]
class IsInModel(torch.nn.Module):
def forward(self, x, test_elements):
return torch.isin(x, test_elements)
@tvm.script.ir_module
class expected:
@R.function
def main(
inp_0: R.Tensor((10, 10), dtype="float32"), inp_1: R.Tensor((8,), dtype="float32")
) -> R.Tensor((10, 10), dtype="bool"):
with R.dataflow():
lv: R.Tensor((10, 10, 1), dtype="float32") = R.expand_dims(inp_0, axis=[-1])
lv1: R.Tensor((8,), dtype="float32") = R.reshape(inp_1, R.shape([8]))
lv2: R.Tensor((10, 10, 8), dtype="bool") = R.equal(lv, lv1)
lv3: R.Tensor((10, 10), dtype="bool") = R.sum(lv2, axis=[-1], keepdims=False)
lv4: R.Tensor((10, 10), dtype="bool") = R.greater(lv3, R.const(0.0, "float32"))
gv: R.Tensor((10, 10), dtype="bool") = lv4
R.output(gv)
return gv
verify_model(IsInModel(), input_info, {}, expected)
def test_div_mode():
input_info = [([64, 64], "float32"), ([64, 64], "float32")]
# Case 1: Basic division (no rounding mode)
class DivModel(torch.nn.Module):
def forward(self, x, y):
return torch.div(x, y)
@tvm.script.ir_module
class expected_div:
@R.function
def main(
inp_0: R.Tensor((64, 64), dtype="float32"), inp_1: R.Tensor((64, 64), dtype="float32")
) -> R.Tensor((64, 64), dtype="float32"):
with R.dataflow():
lv: R.Tensor((64, 64), dtype="float32") = R.divide(inp_0, inp_1)
gv: R.Tensor((64, 64), dtype="float32") = lv
R.output(gv)
return gv
# Case 2: Division with trunc rounding
class DivTruncModel(torch.nn.Module):
def forward(self, x, y):
return torch.div(x, y, rounding_mode="trunc")
@tvm.script.ir_module
class expected_div_trunc:
@R.function
def main(
inp_0: R.Tensor((64, 64), dtype="float32"), inp_1: R.Tensor((64, 64), dtype="float32")
) -> R.Tensor((64, 64), dtype="float32"):
with R.dataflow():
lv: R.Tensor((64, 64), dtype="float32") = R.divide(inp_0, inp_1)
lv1: R.Tensor((64, 64), dtype="float32") = R.trunc(lv)
gv: R.Tensor((64, 64), dtype="float32") = lv1
R.output(gv)
return gv
# Case 3: Division with floor rounding
class DivFloorModel(torch.nn.Module):
def forward(self, x, y):
return torch.div(x, y, rounding_mode="floor")
@tvm.script.ir_module
class expected_div_floor:
@R.function
def main(
inp_0: R.Tensor((64, 64), dtype="float32"), inp_1: R.Tensor((64, 64), dtype="float32")
) -> R.Tensor((64, 64), dtype="float32"):
with R.dataflow():
lv: R.Tensor((64, 64), dtype="float32") = R.floor_divide(inp_0, inp_1)
gv: R.Tensor((64, 64), dtype="float32") = lv
R.output(gv)
return gv
verify_model(DivModel(), input_info, {}, expected_div)
verify_model(DivTruncModel(), input_info, {}, expected_div_trunc)
verify_model(DivFloorModel(), input_info, {}, expected_div_floor)
def test_size():
input_info = [([1, 3, 10, 10], "float32")]
class Size(Module):
def forward(self, input):
return input.size()
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Shape([1, 3, 10, 10]):
# block 0
with R.dataflow():
gv: R.Shape([1, 3, 10, 10]) = R.shape([1, 3, 10, 10])
R.output(gv)
return gv
verify_model(Size(), input_info, {}, expected1)
def test_squeeze():
input_info = [([3, 1, 4, 1], "float32")]
class Squeeze1(Module):
def forward(self, input):
return input.squeeze(1)
@tvm.script.ir_module
class Expected1:
@R.function
def main(inp_0: R.Tensor((3, 1, 4, 1), dtype="float32")) -> R.Tensor(
(3, 4, 1), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((3, 4, 1), dtype="float32") = R.squeeze(inp_0, axis=[1])
gv: R.Tensor((3, 4, 1), dtype="float32") = lv
R.output(gv)
return gv
class Squeeze2(Module):
def forward(self, input):
return input.squeeze()
@tvm.script.ir_module
class Expected2:
@R.function
def main(inp_0: R.Tensor((3, 1, 4, 1), dtype="float32")) -> R.Tensor(
(3, 4), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((3, 4), dtype="float32") = R.squeeze(inp_0, axis=None)
gv: R.Tensor((3, 4), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Squeeze1(), input_info, {}, Expected1)
verify_model(Squeeze2(), input_info, {}, Expected2)
def test_unsqueeze():
input_info = [([1, 3, 10, 10], "float32")]
class Unsqueeze1(Module):
def forward(self, input):
return input.unsqueeze(1)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 1, 3, 10, 10), dtype="float32") = R.expand_dims(input_1, 1)
gv: R.Tensor((1, 1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
class Unsqueeze2(Module):
def forward(self, input):
return input.unsqueeze(-1)
@tvm.script.ir_module
class expected2:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10, 1), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10, 1), dtype="float32") = R.expand_dims(input_1, -1)
gv: R.Tensor((1, 3, 10, 10, 1), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Unsqueeze1(), input_info, {}, expected1)
verify_model(Unsqueeze2(), input_info, {}, expected2)
def test_getattr():
input_info = [([1, 3, 10, 10], "float32")]
class GetAttr1(Module):
def forward(self, input):
return input.shape
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Shape([1, 3, 10, 10]):
# block 0
with R.dataflow():
gv: R.Shape([1, 3, 10, 10]) = R.shape([1, 3, 10, 10])
R.output(gv)
return gv
verify_model(GetAttr1(), input_info, {}, expected1)
def test_getitem():
class Slice1(Module):
def forward(self, x):
return x[0, 1::2, :, :3]
@tvm.script.ir_module
class expected1:
@R.function
def main(x: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 1, 10, 3), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 1, 10, 3), dtype="float32") = R.strided_slice(
x,
axes=[0, 1, 2, 3],
begin=[0, 1, 0, 0],
end=[1, T.int64(3), T.int64(10), 3],
strides=[1, 2, 1, 1],
)
lv1: R.Tensor((1, 1, 10, 3), dtype="float32") = R.reshape(lv, (1, 1, 10, 3))
gv: R.Tensor((1, 1, 10, 3), dtype="float32") = lv1
R.output(gv)
return gv
class Slice2(Module):
def forward(self, x):
return x[:, None, None, :, None]
@I.ir_module
class expected2:
@R.function
def main(inp_0: R.Tensor((8, 16), dtype="float32")) -> R.Tensor(
(8, 1, 1, 16, 1), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((8, 16), dtype="float32") = R.strided_slice(
inp_0, axes=[0, 1], begin=[0, 0], end=[8, 16], strides=[1, 1]
)
lv1: R.Tensor((8, 1, 1, 16, 1), dtype="float32") = R.reshape(
lv, R.shape([8, 1, 1, 16, 1])
)
gv: R.Tensor((8, 1, 1, 16, 1), dtype="float32") = lv1
R.output(gv)
return gv
verify_model(Slice1(), [([1, 3, 10, 10], "float32")], {}, expected1)
verify_model(Slice2(), [([8, 16], "float32")], {}, expected2)
operator_basic_unary = [
(torch.abs, R.abs),
(torch.acos, R.acos),
(torch.acosh, R.acosh),
(torch.asin, R.asin),
(torch.asinh, R.asinh),
(torch.atan, R.atan),
(torch.atanh, R.atanh),
(torch.bitwise_not, R.bitwise_not),
(torch.ceil, R.ceil),
(torch.cos, R.cos),
(torch.cosh, R.cosh),
(torch.erf, R.erf),
(torch.exp, R.exp),
(torch.floor, R.floor),
(torch.log, R.log),
(torch.neg, R.negative),
(torch.round, R.round),
(torch.rsqrt, R.rsqrt),
(torch.sin, R.sin),
(torch.sinh, R.sinh),
(torch.sign, R.sign),
(torch.sqrt, R.sqrt),
(torch.square, R.square),
(torch.tan, R.tan),
]
@pytest.mark.parametrize("pytorch_op, relax_op", operator_basic_unary)
def test_basic_unary_ops(pytorch_op, relax_op):
input_info = [([1, 3, 10, 10], "float32")]
class Unary(Module):
def forward(self, input):
return pytorch_op(input)
@tvm.script.ir_module
class expected_unary:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(input_1)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Unary(), input_info, {}, expected_unary)
def test_sqrt_integer_input_fx():
input_info = [([1, 4], "int64")]
class SqrtIntModel(Module):
def forward(self, input):
return torch.sqrt(input)
@tvm.script.ir_module
class expected:
@R.function
def main(input_1: R.Tensor((1, 4), dtype="int64")) -> R.Tensor((1, 4), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 4), dtype="float32") = R.astype(input_1, dtype="float32")
lv1: R.Tensor((1, 4), dtype="float32") = R.sqrt(lv)
gv: R.Tensor((1, 4), dtype="float32") = lv1
R.output(gv)
return gv
verify_model(SqrtIntModel(), input_info, {}, expected)
operator_bool_unary = [
(torch.isnan, R.isnan),
(torch.isinf, R.isinf),
(torch.isfinite, R.isfinite),
]
@pytest.mark.parametrize("pytorch_op, relax_op", operator_bool_unary)
def test_bool_unary_ops(pytorch_op, relax_op):
input_info = [([1, 3, 10, 10], "float32")]
class Unary(Module):
def forward(self, input):
return pytorch_op(input)
@tvm.script.ir_module
class expected_unary:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="bool"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = relax_op(input_1)
gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv
R.output(gv)
return gv
verify_model(Unary(), input_info, {}, expected_unary)
def test_extended_unary_ops():
input_info = [([1, 3, 10, 10], "float32")]
# celu
class Celu1(Module):
def __init__(self):
super().__init__()
self.celu = torch.nn.CELU()
def forward(self, input):
return self.celu(input)
class Celu2(Module):
def forward(self, input):
return torch.nn.functional.celu(input)
# alpha * min(0, exp(x / alpha) - 1) + max(0, x)
@tvm.script.ir_module
class expected_celu:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(input_1)
lv_div: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
lv, R.const(1.0, "float32")
)
lv_sub: R.Tensor((1, 3, 10, 10), dtype="float32") = R.subtract(
lv_div, R.const(1.0, "float32")
)
lv_min: R.Tensor((1, 3, 10, 10), dtype="float32") = R.minimum(
R.const(0.0, "float32"), lv_sub
)
lv_scaled: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(
R.const(1.0, "float32"), lv_min
)
lv_relu_x: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input_1)
lv_celu: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv_scaled, lv_relu_x)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv_celu
R.output(gv)
return gv
verify_model(Celu1(), input_info, {}, expected_celu)
verify_model(Celu2(), input_info, {}, expected_celu)
# clamp
class Clamp(Module):
def forward(self, input):
return torch.clamp(input, min=0.1, max=0.5)
@tvm.script.ir_module
class expected_clamp:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(input_1, 0.1, 0.5)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Clamp(), input_info, {}, expected_clamp)
class ClampMinOnly(Module):
def forward(self, input):
return torch.clamp(input, min=0.5, max=None)
@tvm.script.ir_module
class expected_clamp_min_only:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(input_1, 0.5, math.inf)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(ClampMinOnly(), input_info, {}, expected_clamp_min_only)
class ClampTensors(Module):
def forward(self, input):
return torch.clamp(input, min=input, max=input)
@tvm.script.ir_module
class expected_clamp_tensors:
@R.function
def main(
inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.broadcast_to(
inp_0, R.shape([1, 3, 10, 10])
)
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.maximum(inp_0, lv)
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.broadcast_to(
inp_0, R.shape([1, 3, 10, 10])
)
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.minimum(lv1, lv2)
lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
lv3, R.prim_value(T.float64("-inf")), R.prim_value(T.float64("inf"))
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv4
R.output(gv)
return gv
verify_model(ClampTensors(), input_info, {}, expected_clamp_tensors)
# dropout
class Dropout1(Module):
def __init__(self):
super().__init__()
self.dropout = torch.nn.Dropout(0.5)
def forward(self, input):
return self.dropout(input)
class Dropout2(Module):
def forward(self, input):
return torch.dropout(input, 0.5, train=True)
@tvm.script.ir_module
class expected_dropout:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = input_1
R.output(gv)
return gv
verify_model(Dropout1(), input_info, {}, expected_dropout)
verify_model(Dropout2(), input_info, {}, expected_dropout)
# elu
class Elu(Module):
def __init__(self):
super().__init__()
self.elu = torch.nn.ELU()
def forward(self, input):
return self.elu(input)
class Elu2(Module):
def forward(self, input):
return torch.nn.functional.elu(input)
@tvm.script.ir_module
class expected_elu:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv_exp: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(input_1)
lv_one_minus_exp: R.Tensor((1, 3, 10, 10), dtype="float32") = R.subtract(
R.const(1.0, dtype="float32"), lv_exp
)
lv_relu_one_minus_exp: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(
lv_one_minus_exp
)
lv_scaled: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(
R.const(-1.0, dtype="float32"), lv_relu_one_minus_exp
)
lv_relu_x: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input_1)
lv_elu: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv_scaled, lv_relu_x)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv_elu
R.output(gv)
return gv
verify_model(Elu(), input_info, {}, expected_elu)
verify_model(Elu2(), input_info, {}, expected_elu)
# gelu
class Gelu(Module):
def __init__(self):
super().__init__()
self.gelu = torch.nn.GELU()
def forward(self, input):
return self.gelu(input)
class Gelu2(Module):
def forward(self, input):
return torch.nn.functional.gelu(input)
@tvm.script.ir_module
class expected_gelu:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.gelu(input_1)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Gelu(), input_info, {}, expected_gelu)
verify_model(Gelu2(), input_info, {}, expected_gelu)
# hardsigmoid
class Hardsigmoid(torch.nn.Module):
def __init__(self):
super().__init__()
self.hs = torch.nn.Hardsigmoid()
def forward(self, input):
return self.hs(input)
class Hardsigmoid2(torch.nn.Module):
def forward(self, input):
return torch.nn.functional.hardsigmoid(input)
@tvm.script.ir_module
class expected_hardsigmoid:
@R.function
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(inp_0, R.const(3, "float32"))
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(lv, 0, 6)
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
lv1, R.const(6, "float32")
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv2
R.output(gv)
return gv
verify_model(Hardsigmoid(), input_info, {}, expected_hardsigmoid)
verify_model(Hardsigmoid2(), input_info, {}, expected_hardsigmoid)
# hardswish
class Hardswish(torch.nn.Module):
def __init__(self):
super().__init__()
self.hs = torch.nn.Hardswish()
def forward(self, input):
return self.hs(input)
class Hardswish2(torch.nn.Module):
def forward(self, input):
return torch.nn.functional.hardswish(input)
@tvm.script.ir_module
class expected_hardswish:
@R.function
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(inp_0, R.const(3, "float32"))
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(lv, 0, 6)
lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
lv1, R.const(6, "float32")
)
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply(inp_0, lv2)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv3
R.output(gv)
return gv
verify_model(Hardswish(), input_info, {}, expected_hardswish)
verify_model(Hardswish2(), input_info, {}, expected_hardswish)
# hardtanh
class Hardtanh(torch.nn.Module):
def __init__(self):
super().__init__()
self.ht = torch.nn.Hardtanh()
def forward(self, input):
return self.ht(input)
class Hardtanh2(torch.nn.Module):
def forward(self, input):
return torch.nn.functional.hardtanh(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
inp_0, R.prim_value(T.float64(-1.0)), R.prim_value(T.float64(1.0))
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Hardtanh(), input_info, {}, expected1)
verify_model(Hardtanh2(), input_info, {}, expected1)
# leaky_relu
test_leakyrelu()
# softplus
test_softplus()
# prelu
test_prelu()
# log2
class Log2(Module):
def forward(self, x):
return torch.log2(x)
@tvm.script.ir_module
class Expected_log2:
@R.function
def main(
inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(inp_0)
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
lv, R.const(0.6931471805599453, "float32")
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv1
R.output(gv)
return gv
verify_model(Log2(), input_info, {}, Expected_log2)
# log10
class Log10(Module):
def forward(self, x):
return torch.log10(x)
@tvm.script.ir_module
class Expected_log10:
@R.function
def main(
inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(inp_0)
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
lv, R.const(2.302585092994046, "float32")
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv1
R.output(gv)
return gv
verify_model(Log10(), input_info, {}, Expected_log10)
# log1p
class Log1p(Module):
def forward(self, x):
return torch.log1p(x)
@tvm.script.ir_module
class Expected_log1p:
@R.function
def main(
inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(
R.add(inp_0, R.const(1, "float32"))
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Log1p(), input_info, {}, Expected_log1p)
# logical_not
class LogicalNot(Module):
def forward(self, input):
return torch.logical_not(input)
@tvm.script.ir_module
class expected_logical_not:
@R.function
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="bool"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(inp_0, dtype="bool")
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_not(lv)
gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv1
R.output(gv)
return gv
verify_model(LogicalNot(), input_info, {}, expected_logical_not)
# log_softmax
class LogSoftmax(Module):
def __init__(self):
super().__init__()
self.lsm = torch.nn.LogSoftmax(dim=1)
def forward(self, input):
return self.lsm(input)
class LogSoftmax2(Module):
def forward(self, input):
return torch.nn.functional.log_softmax(input, dim=1)
@tvm.script.ir_module
class expected_log_softmax:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.log_softmax(input_1, axis=1)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(LogSoftmax(), input_info, {}, expected_log_softmax)
verify_model(LogSoftmax2(), input_info, {}, expected_log_softmax)
# reciprocal
class Reciprocal(Module):
def forward(self, input):
return torch.reciprocal(input)
@tvm.script.ir_module
class expected_reciprocal:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
R.const(1.0, "float32"), input_1
)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Reciprocal(), input_info, {}, expected_reciprocal)
# relu
class ReLU0(Module):
def __init__(self):
super().__init__()
self.relu = torch.nn.ReLU()
def forward(self, input):
return self.relu(input)
class ReLU1(Module):
def forward(self, input):
return torch.nn.functional.relu(input)
@tvm.script.ir_module
class expected_relu:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu(input_1)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(ReLU0(), input_info, {}, expected_relu)
verify_model(ReLU1(), input_info, {}, expected_relu)
# relu6
class ReLU6_1(torch.nn.Module):
def __init__(self):
super().__init__()
self.relu6 = torch.nn.ReLU6()
def forward(self, x):
return self.relu6(x)
class ReLU6_2(torch.nn.Module):
def forward(self, x):
return torch.nn.functional.relu6(x)
@tvm.script.ir_module
class expected:
@R.function
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.relu6(inp_0)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(ReLU6_1(), input_info, {}, expected)
verify_model(ReLU6_2(), input_info, {}, expected)
# selu
class Selu1(Module):
def __init__(self):
super().__init__()
self.selu = torch.nn.SELU()
def forward(self, input):
return self.selu(input)
class Selu2(Module):
def forward(self, input):
return torch.nn.functional.selu(input)
@tvm.script.ir_module
class expected_selu:
@R.function
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.selu(inp_0)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Selu1(), input_info, {}, expected_selu)
verify_model(Selu2(), input_info, {}, expected_selu)
# sigmoid
class Sigmoid(Module):
def __init__(self):
super().__init__()
self.sigmoid = torch.nn.Sigmoid()
def forward(self, input):
return self.sigmoid(input)
class Sigmoid2(Module):
def forward(self, input):
return torch.sigmoid(input)
@tvm.script.ir_module
class expected_sigmoid:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.sigmoid(input_1)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Sigmoid(), input_info, {}, expected_sigmoid)
verify_model(Sigmoid2(), input_info, {}, expected_sigmoid)
# silu
class SiLU(Module):
def __init__(self):
super().__init__()
self.silu = torch.nn.SiLU()
def forward(self, input):
return self.silu(input)
class SiLU2(Module):
def forward(self, input):
return torch.nn.functional.silu(input)
@tvm.script.ir_module
class expected_silu:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.silu(input_1)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(SiLU(), input_info, {}, expected_silu)
verify_model(SiLU2(), input_info, {}, expected_silu)
# softmax
class Softmax(Module):
def __init__(self):
super().__init__()
self.sm = torch.nn.Softmax(dim=1)
def forward(self, input):
return self.sm(input)
class Softmax2(Module):
def forward(self, input):
return torch.nn.functional.softmax(input, dim=1)
@tvm.script.ir_module
class expected_softmax:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.nn.softmax(input_1, axis=1)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Softmax(), input_info, {}, expected_softmax)
verify_model(Softmax2(), input_info, {}, expected_softmax)
# tanh
class Tanh(Module):
def __init__(self):
super().__init__()
self.tanh = torch.nn.Tanh()
def forward(self, input):
return self.tanh(input)
class Tanh2(Module):
def forward(self, input):
return torch.tanh(input)
@tvm.script.ir_module
class expected_tanh:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.tanh(input_1)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Tanh(), input_info, {}, expected_tanh)
verify_model(Tanh2(), input_info, {}, expected_tanh)
# tril
class Tril(Module):
def forward(self, input):
return torch.tril(input, 1)
class InplaceTril(Module):
def forward(self, input):
input.tril_(1)
return input
@tvm.script.ir_module
class expected_tril:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.tril(input_1, 1)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Tril(), input_info, {}, expected_tril)
verify_model(InplaceTril(), input_info, {}, expected_tril)
# triu
class Triu(Module):
def forward(self, input):
return torch.triu(input, 1)
class InplaceTriu(Module):
def forward(self, input):
input.triu_(1)
return input
@tvm.script.ir_module
class expected_triu:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.triu(input_1, 1)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Triu(), input_info, {}, expected_triu)
verify_model(InplaceTriu(), input_info, {}, expected_triu)
# trunc
class Trunc(torch.nn.Module):
def forward(self, input):
return torch.trunc(input)
@tvm.script.ir_module
class expected_trunc:
@R.function
def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 10, 10), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.trunc(inp_0)
gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Trunc(), input_info, {}, expected_trunc)
def test_logical_and():
input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")]
class LogicalAnd(Module):
def forward(self, lhs, rhs):
return torch.logical_and(lhs, rhs)
@tvm.script.ir_module
class expected:
@R.function
def main(
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="bool"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool")
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool")
lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_and(lv, lv1)
gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2
R.output(gv)
return gv
verify_model(LogicalAnd(), input_info, {}, expected)
def test_logical_or():
input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")]
class LogicalOr(Module):
def forward(self, lhs, rhs):
return torch.logical_or(lhs, rhs)
@tvm.script.ir_module
class expected:
@R.function
def main(
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="bool"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool")
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool")
lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_or(lv, lv1)
gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2
R.output(gv)
return gv
verify_model(LogicalOr(), input_info, {}, expected)
def test_logical_xor():
input_info = [([1, 3, 10, 10], "float32"), ([1, 3, 10, 10], "float32")]
class LogicalXor(Module):
def forward(self, lhs, rhs):
return torch.logical_xor(lhs, rhs)
@tvm.script.ir_module
class expected:
@R.function
def main(
lhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
rhs: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((1, 3, 10, 10), dtype="bool"):
with R.dataflow():
lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(lhs, dtype="bool")
lv1: R.Tensor((1, 3, 10, 10), dtype="bool") = R.astype(rhs, dtype="bool")
lv2: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_xor(lv, lv1)
gv: R.Tensor((1, 3, 10, 10), dtype="bool") = lv2
R.output(gv)
return gv
verify_model(LogicalXor(), input_info, {}, expected)
def test_pow_integer():
input_info = [([4], "int64")]
class Pow(Module):
def forward(self, input):
return input.pow(4)
@tvm.script.ir_module
class expected:
@R.function
def main(inp_0: R.Tensor((4,), dtype="int64")) -> R.Tensor((4,), dtype="int64"):
with R.dataflow():
lv: R.Tensor((4,), dtype="int64") = R.multiply(inp_0, inp_0)
lv1: R.Tensor((4,), dtype="int64") = R.multiply(lv, inp_0)
lv2: R.Tensor((4,), dtype="int64") = R.multiply(lv1, inp_0)
gv: R.Tensor((4,), dtype="int64") = lv2
R.output(gv)
return gv
verify_model(Pow(), input_info, {}, expected)
def test_interpolate():
input_info = [([1, 3, 10, 10], "float32")]
class Interpolate(Module):
def forward(self, input):
return torch.nn.functional.interpolate(input, (5, 5))
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 5, 5), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 5, 5), dtype="float32") = R.image.resize2d(
input_1,
(5, 5),
roi=[0.000000, 0.000000, 0.000000, 0.000000],
layout="NCHW",
method="nearest_neighbor",
coordinate_transformation_mode="asymmetric",
rounding_method="round",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 3, 5, 5), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Interpolate(), input_info, {}, expected1)
class Interpolate2(Module):
def forward(self, input):
return torch.nn.functional.interpolate(
input,
size=None,
scale_factor=2.0,
mode="bilinear",
align_corners=False,
)
@tvm.script.ir_module
class expected2:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 20, 20), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 20, 20), dtype="float32") = R.image.resize2d(
input_1,
(20, 20),
roi=[0.000000, 0.000000, 0.000000, 0.000000],
layout="NCHW",
method="linear",
coordinate_transformation_mode="half_pixel",
rounding_method="round",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 3, 20, 20), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Interpolate2(), input_info, {}, expected2)
class Interpolate3(Module):
def forward(self, input):
return torch.nn.functional.interpolate(
input,
size=None,
scale_factor=(2.0, 1.0),
mode="bilinear",
align_corners=False,
)
@tvm.script.ir_module
class expected3:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 20, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 20, 10), dtype="float32") = R.image.resize2d(
input_1,
(20, 10),
roi=[0.000000, 0.000000, 0.000000, 0.000000],
layout="NCHW",
method="linear",
coordinate_transformation_mode="half_pixel",
rounding_method="round",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 3, 20, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Interpolate3(), input_info, {}, expected3)
class Interpolate4(Module):
def forward(self, input):
return torch.nn.functional.interpolate(
input,
size=None,
scale_factor=(2.0, 1.0),
mode="bicubic",
align_corners=False,
)
@tvm.script.ir_module
class expected4:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 20, 10), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 20, 10), dtype="float32") = R.image.resize2d(
input_1,
(20, 10),
roi=[0.000000, 0.000000, 0.000000, 0.000000],
layout="NCHW",
method="cubic",
coordinate_transformation_mode="half_pixel",
rounding_method="round",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 3, 20, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Interpolate4(), input_info, {}, expected4)
input_info_5d = [([1, 3, 4, 10, 10], "float32")]
class Interpolate5(Module):
def forward(self, input):
return torch.nn.functional.interpolate(
input,
size=None,
scale_factor=(2.0, 2.0, 2.0),
mode="trilinear",
align_corners=False,
)
@tvm.script.ir_module
class expected5:
@R.function
def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 8, 20, 20), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 8, 20, 20), dtype="float32") = R.image.resize3d(
input_5,
(8, 20, 20),
roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
layout="NCDHW",
method="linear",
coordinate_transformation_mode="half_pixel",
rounding_method="",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 3, 8, 20, 20), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Interpolate5(), input_info_5d, {}, expected5)
class Interpolate6(Module):
def forward(self, input):
return torch.nn.functional.interpolate(
input,
size=None,
scale_factor=(2.0, 4.0, 4.0),
mode="trilinear",
align_corners=False,
)
@tvm.script.ir_module
class expected6:
@R.function
def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 8, 40, 40), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d(
input_5,
(8, 40, 40),
roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
layout="NCDHW",
method="linear",
coordinate_transformation_mode="half_pixel",
rounding_method="",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Interpolate6(), input_info_5d, {}, expected6)
class Interpolate7(Module):
def forward(self, input):
return torch.nn.functional.interpolate(
input,
size=(8, 40, 40),
mode="trilinear",
align_corners=False,
)
@tvm.script.ir_module
class expected7:
@R.function
def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 8, 40, 40), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d(
input_5,
(8, 40, 40),
roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
layout="NCDHW",
method="linear",
coordinate_transformation_mode="half_pixel",
rounding_method="",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Interpolate7(), input_info_5d, {}, expected7)
class Interpolate8(Module):
def forward(self, input):
return torch.nn.functional.interpolate(
input,
size=(8, 40, 40),
mode="trilinear",
align_corners=True,
)
@tvm.script.ir_module
class expected8:
@R.function
def main(input_5: R.Tensor((1, 3, 4, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 8, 40, 40), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = R.image.resize3d(
input_5,
(8, 40, 40),
roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
layout="NCDHW",
method="linear",
coordinate_transformation_mode="align_corners",
rounding_method="",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 3, 8, 40, 40), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Interpolate8(), input_info_5d, {}, expected8)
def test_interpolate_nhwc_layout():
# First verify backward compatibility - default should still be NCHW
input_info_nchw = [([1, 3, 10, 10], "float32")]
class InterpolateDefault(Module):
def forward(self, input):
return torch.nn.functional.interpolate(input, (5, 5))
@tvm.script.ir_module
class expected_default_nchw:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tensor(
(1, 3, 5, 5), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 3, 5, 5), dtype="float32") = R.image.resize2d(
input_1,
(5, 5),
roi=[0.000000, 0.000000, 0.000000, 0.000000],
layout="NCHW",
method="nearest_neighbor",
coordinate_transformation_mode="asymmetric",
rounding_method="round",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 3, 5, 5), dtype="float32") = lv
R.output(gv)
return gv
# Verify default behavior (no default_image_layout parameter) uses NCHW
graph_model_default = fx.symbolic_trace(InterpolateDefault())
with torch.no_grad():
mod_default = from_fx(graph_model_default, input_info_nchw)
tvm.ir.assert_structural_equal(mod_default, expected_default_nchw)
# Now test NHWC layout
input_info = [([1, 10, 10, 3], "float32")]
class InterpolateNHWC(Module):
def forward(self, input):
return torch.nn.functional.interpolate(input, (5, 5))
@tvm.script.ir_module
class expected_nhwc:
@R.function
def main(input_1: R.Tensor((1, 10, 10, 3), dtype="float32")) -> R.Tensor(
(1, 5, 5, 3), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 5, 5, 3), dtype="float32") = R.image.resize2d(
input_1,
(5, 5),
roi=[0.000000, 0.000000, 0.000000, 0.000000],
layout="NHWC",
method="nearest_neighbor",
coordinate_transformation_mode="asymmetric",
rounding_method="round",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 5, 5, 3), dtype="float32") = lv
R.output(gv)
return gv
# Test with NHWC layout
graph_model = fx.symbolic_trace(InterpolateNHWC())
with torch.no_grad():
mod = from_fx(graph_model, input_info, default_image_layout="NHWC")
tvm.ir.assert_structural_equal(mod, expected_nhwc)
# Test with bilinear interpolation and NHWC layout
class InterpolateNHWC2(Module):
def forward(self, input):
return torch.nn.functional.interpolate(
input, size=None, scale_factor=2.0, mode="bilinear", align_corners=False
)
@tvm.script.ir_module
class expected_nhwc2:
@R.function
def main(input_1: R.Tensor((1, 10, 10, 3), dtype="float32")) -> R.Tensor(
(1, 20, 20, 3), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 20, 20, 3), dtype="float32") = R.image.resize2d(
input_1,
(20, 20),
roi=[0.000000, 0.000000, 0.000000, 0.000000],
layout="NHWC",
method="linear",
coordinate_transformation_mode="half_pixel",
rounding_method="round",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 20, 20, 3), dtype="float32") = lv
R.output(gv)
return gv
graph_model2 = fx.symbolic_trace(InterpolateNHWC2())
with torch.no_grad():
mod2 = from_fx(graph_model2, input_info, default_image_layout="NHWC")
tvm.ir.assert_structural_equal(mod2, expected_nhwc2)
input_info_5d = [([1, 4, 10, 10, 3], "float32")]
class InterpolateNHWC3(Module):
def forward(self, input):
return torch.nn.functional.interpolate(
input,
size=None,
scale_factor=(2.0, 4.0, 4.0),
mode="trilinear",
align_corners=False,
)
@tvm.script.ir_module
class expected_nhwc3:
@R.function
def main(input_5: R.Tensor((1, 4, 10, 10, 3), dtype="float32")) -> R.Tensor(
(1, 8, 40, 40, 3), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = R.image.resize3d(
input_5,
(8, 40, 40),
roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
layout="NDHWC",
method="linear",
coordinate_transformation_mode="half_pixel",
rounding_method="",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = lv
R.output(gv)
return gv
graph_model3 = fx.symbolic_trace(InterpolateNHWC3())
with torch.no_grad():
mod3 = from_fx(graph_model3, input_info_5d, default_image_layout="NDHWC")
tvm.ir.assert_structural_equal(mod3, expected_nhwc3)
class InterpolateNHWC4(Module):
def forward(self, input):
return torch.nn.functional.interpolate(
input,
size=None,
scale_factor=(2.0, 4.0, 4.0),
mode="trilinear",
align_corners=True,
)
@tvm.script.ir_module
class expected_nhwc4:
@R.function
def main(input_5: R.Tensor((1, 4, 10, 10, 3), dtype="float32")) -> R.Tensor(
(1, 8, 40, 40, 3), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = R.image.resize3d(
input_5,
(8, 40, 40),
roi=[0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000],
layout="NDHWC",
method="linear",
coordinate_transformation_mode="align_corners",
rounding_method="",
cubic_alpha=-0.75,
cubic_exclude=0,
extrapolation_value=0,
)
gv: R.Tensor((1, 8, 40, 40, 3), dtype="float32") = lv
R.output(gv)
return gv
graph_model4 = fx.symbolic_trace(InterpolateNHWC4())
with torch.no_grad():
mod4 = from_fx(graph_model4, input_info_5d, default_image_layout="NDHWC")
tvm.ir.assert_structural_equal(mod4, expected_nhwc4)
def test_addmm():
input_info = [
([10, 10], "float32"),
([10, 10], "float32"),
([10, 10], "float32"),
]
class Addmm1(Module):
def __init__(self):
super().__init__()
def forward(self, x1, x2, x3):
return torch.addmm(x1, x2, x3)
class Addmm2(Module):
def __init__(self):
super().__init__()
def forward(self, x1, x2, x3):
return torch.addmm(x1, x2, x3, beta=0.8, alpha=0.5)
@tvm.script.ir_module
class expected1:
@R.function
def main(
x1: R.Tensor((10, 10), dtype="float32"),
x2: R.Tensor((10, 10), dtype="float32"),
x3: R.Tensor((10, 10), dtype="float32"),
) -> R.Tensor((10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((10, 10), dtype="float32") = R.matmul(x2, x3, out_dtype="float32")
lv1: R.Tensor((10, 10), dtype="float32") = R.add(x1, lv)
gv: R.Tensor((10, 10), dtype="float32") = lv1
R.output(gv)
return gv
@tvm.script.ir_module
class expected2:
@R.function
def main(
x1: R.Tensor((10, 10), dtype="float32"),
x2: R.Tensor((10, 10), dtype="float32"),
x3: R.Tensor((10, 10), dtype="float32"),
) -> R.Tensor((10, 10), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((10, 10), dtype="float32") = R.matmul(x2, x3, out_dtype="float32")
lv1: R.Tensor((10, 10), dtype="float32") = R.multiply(lv, R.const(0.5, "float32"))
lv2: R.Tensor((10, 10), dtype="float32") = R.multiply(x1, R.const(0.8, "float32"))
lv3: R.Tensor((10, 10), dtype="float32") = R.add(lv2, lv1)
gv: R.Tensor((10, 10), dtype="float32") = lv3
R.output(gv)
return gv
verify_model(Addmm1(), input_info, {}, expected1)
verify_model(Addmm2(), input_info, {}, expected2)
def test_split():
input_info = [([1, 3, 10, 10], "float32")]
class Split1(Module):
def forward(self, input):
return torch.split(input, 1, dim=1)
class Split2(Module):
def forward(self, input):
return torch.split(input, [1, 2], dim=1)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
):
# block 0
with R.dataflow():
lv: R.Tuple(
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
) = R.split(input_1, indices_or_sections=3, axis=1)
gv: R.Tuple(
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
) = lv
R.output(gv)
return gv
@tvm.script.ir_module
class expected2:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
R.Tensor((1, 1, 10, 10), dtype="float32"), R.Tensor((1, 2, 10, 10), dtype="float32")
):
# block 0
with R.dataflow():
lv: R.Tuple(
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 2, 10, 10), dtype="float32"),
) = R.split(input_1, indices_or_sections=[1], axis=1)
gv: R.Tuple(
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 2, 10, 10), dtype="float32"),
) = lv
R.output(gv)
return gv
verify_model(Split1(), input_info, {}, expected1)
verify_model(Split2(), input_info, {}, expected2)
def test_unbind():
input_info = [([3, 3, 10, 10], "float32")]
class Unbind1(Module):
def forward(self, data):
return torch.unbind(data)
class Unbind2(Module):
def forward(self, data):
return torch.unbind(data, dim=1)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((3, 3, 10, 10), dtype="float32")) -> R.Tuple(
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
):
# block 0
with R.dataflow():
lv: R.Tuple(
R.Tensor((1, 3, 10, 10), dtype="float32"),
R.Tensor((1, 3, 10, 10), dtype="float32"),
R.Tensor((1, 3, 10, 10), dtype="float32"),
) = R.split(input_1, indices_or_sections=3, axis=0)
lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[0]
lv2: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv1, axis=[0])
lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[1]
lv4: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv3, axis=[0])
lv5: R.Tensor((1, 3, 10, 10), dtype="float32") = lv[2]
lv6: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv5, axis=[0])
lv7: R.Tuple(
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
) = (lv2, lv4, lv6)
gv: R.Tuple(
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
) = lv7
R.output(gv)
return gv
@tvm.script.ir_module
class expected2:
@R.function
def main(input_1: R.Tensor((3, 3, 10, 10), dtype="float32")) -> R.Tuple(
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
):
# block 0
with R.dataflow():
lv: R.Tuple(
R.Tensor((3, 1, 10, 10), dtype="float32"),
R.Tensor((3, 1, 10, 10), dtype="float32"),
R.Tensor((3, 1, 10, 10), dtype="float32"),
) = R.split(input_1, indices_or_sections=3, axis=1)
lv1: R.Tensor((3, 1, 10, 10), dtype="float32") = lv[0]
lv2: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv1, axis=[1])
lv3: R.Tensor((3, 1, 10, 10), dtype="float32") = lv[1]
lv4: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv3, axis=[1])
lv5: R.Tensor((3, 1, 10, 10), dtype="float32") = lv[2]
lv6: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv5, axis=[1])
lv7: R.Tuple(
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
) = (lv2, lv4, lv6)
gv: R.Tuple(
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
R.Tensor((3, 10, 10), dtype="float32"),
) = lv7
R.output(gv)
return gv
verify_model(Unbind1(), input_info, {}, expected1)
verify_model(Unbind2(), input_info, {}, expected2)
def test_cumsum():
input_info = [([1, 2, 3, 4], "float32")]
class Cumsum(Module):
def forward(self, input):
return torch.cumsum(input, dim=1, dtype=torch.int32)
@tvm.script.ir_module
class expected1:
@R.function
def main(input_1: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
(1, 2, 3, 4), dtype="int32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 2, 3, 4), dtype="int32") = R.cumsum(input_1, axis=1, dtype="int32")
gv: R.Tensor((1, 2, 3, 4), dtype="int32") = lv
R.output(gv)
return gv
verify_model(Cumsum(), input_info, {}, expected1)
def test_chunk():
input_info = [([1, 3, 10, 10], "float32")]
class Chunk(Module):
def forward(self, input):
return torch.chunk(input, 3, dim=1)
@tvm.script.ir_module
class Expected:
@R.function
def main(input_1: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
):
# block 0
with R.dataflow():
lv: R.Tuple(
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
) = R.split(input_1, indices_or_sections=3, axis=1)
gv: R.Tuple(
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
R.Tensor((1, 1, 10, 10), dtype="float32"),
) = lv
R.output(gv)
return gv
verify_model(Chunk(), input_info, {}, Expected)
def test_inplace_fill():
class InplaceFill(Module):
def forward(self, input):
input.fill_(1.5)
return input
@tvm.script.ir_module
class Expected:
@R.function
def main(inp_0: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((10, 10), dtype="float32") = R.full(
R.shape([10, 10]), R.const(1.5, "float32"), dtype="float32"
)
gv: R.Tensor((10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(InplaceFill(), [([10, 10], "float32")], {}, Expected)
def test_masked_fill_inplace():
class Masked_Fill_Inplace(Module):
def forward(self, input: torch.Tensor, mask: torch.Tensor):
input.masked_fill_(mask, 1.5)
return input
@tvm.script.ir_module
class Expected:
@R.function
def main(
input: R.Tensor((10, 10), dtype="float32"), mask: R.Tensor((10, 10), dtype="bool")
) -> R.Tensor((10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((10, 10), dtype="float32") = R.full_like(
input, R.const(1.5, "float32")
)
lv1: R.Tensor((10, 10), dtype="float32") = R.where(mask, lv, input)
gv: R.Tensor((10, 10), dtype="float32") = lv1
R.output(gv)
return gv
input_info = [((10, 10), "float32"), ((10, 10), "bool")]
verify_model(Masked_Fill_Inplace(), input_info, {}, Expected)
def test_arange():
import numpy as np
torch.set_grad_enabled(False)
torch.random.manual_seed(0)
class Arange(Module):
def forward(self, input):
return torch.arange(0, 20, dtype=torch.int32)
graph_model = fx.symbolic_trace(Arange())
mod = from_fx(graph_model, [([10, 10], "float32")])
assert len(mod["main"].body.blocks) == 1
assert len(mod["main"].body.blocks[0].bindings) == 1
assert isinstance(mod["main"].body.blocks[0].bindings[0].value, relax.Constant)
tvm.testing.assert_allclose(
mod["main"].body.blocks[0].bindings[0].value.data.numpy(),
np.arange(0, 20, dtype="int32"),
)
def test_empty():
class Empty(Module):
def forward(self, input):
return torch.empty((10, 10), dtype=torch.float32)
graph_model = fx.symbolic_trace(Empty())
mod = from_fx(graph_model, [([10, 10], "float32")])
assert len(mod["main"].body.blocks) == 1
assert len(mod["main"].body.blocks[0].bindings) == 1
assert isinstance(mod["main"].body.blocks[0].bindings[0].value, relax.Constant)
assert mod["main"].body.blocks[0].bindings[0].value.data.shape == (10, 10)
assert mod["main"].body.blocks[0].bindings[0].value.data.dtype == "float32"
def test_tensor():
class Empty1(Module):
def forward(self, input):
return torch.tensor(3, dtype=torch.float32)
class Empty2(Module):
def forward(self, input):
return torch.tensor(3)
graph_model1 = fx.symbolic_trace(Empty1())
mod1 = from_fx(graph_model1, [([10, 10], "float32")])
assert len(mod1["main"].body.blocks) == 1
assert len(mod1["main"].body.blocks[0].bindings) == 1
assert isinstance(mod1["main"].body.blocks[0].bindings[0].value, relax.Constant)
assert mod1["main"].body.blocks[0].bindings[0].value.data.shape == ()
assert mod1["main"].body.blocks[0].bindings[0].value.data.dtype == "float32"
graph_model2 = fx.symbolic_trace(Empty2())
mod2 = from_fx(graph_model2, [([10, 10], "float32")])
assert len(mod2["main"].body.blocks) == 1
assert len(mod2["main"].body.blocks[0].bindings) == 1
assert isinstance(mod2["main"].body.blocks[0].bindings[0].value, relax.Constant)
assert mod2["main"].body.blocks[0].bindings[0].value.data.shape == ()
assert mod2["main"].body.blocks[0].bindings[0].value.data.dtype == "int64"
def test_new_ones():
input_info = [([1, 2, 3], "float32")]
class NewOnes(Module):
def forward(self, x):
return x.new_ones(1, 2, 3)
@tvm.script.ir_module
class expected1:
@R.function
def main(x: R.Tensor((1, 2, 3), dtype="float32")) -> R.Tensor((1, 2, 3), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((1, 2, 3), dtype="float32") = R.full(
(1, 2, 3), R.const(1, "float32"), dtype="float32"
)
gv: R.Tensor((1, 2, 3), dtype="float32") = lv
R.output(gv)
return gv
verify_model(NewOnes(), input_info, {}, expected1)
def test_new_zeros():
input_info = [([1, 128, 128], "float32")]
class NewZeros(Module):
def forward(self, x):
return x.new_zeros(1, 128, 128)
@tvm.script.ir_module
class expected:
@R.function
def main(x: R.Tensor((1, 128, 128), dtype="float32")) -> R.Tensor(
(1, 128, 128), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 128, 128), dtype="float32") = R.full(
(1, 128, 128), R.const(0.0, "float32"), dtype="float32"
)
gv: R.Tensor((1, 128, 128), dtype="float32") = lv
R.output(gv)
return gv
verify_model(NewZeros(), input_info, {}, expected)
def test_expand():
input_info = [([1, 2, 3, 4], "float32")]
class Expand1(Module):
def forward(self, x):
return x.expand(4, 2, 3, 4)
class Expand2(Module):
def forward(self, x):
return x.expand(4, -1, -1, 4)
@tvm.script.ir_module
class expected1:
@R.function
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
(4, 2, 3, 4), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((4, 2, 3, 4), dtype="float32") = R.broadcast_to(x, (4, 2, 3, 4))
gv: R.Tensor((4, 2, 3, 4), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Expand1(), input_info, {}, expected1)
verify_model(Expand2(), input_info, {}, expected1)
def test_reduce():
input_info = [([1, 2, 3, 4], "float32")]
# sum
class Sum(Module):
def forward(self, x):
return torch.sum(x, (2, 1))
@tvm.script.ir_module
class expected1:
@R.function
def main(inp_0: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
(1, 4), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 4), dtype="float32") = R.sum(inp_0, axis=[2, 1], keepdims=False)
gv: R.Tensor((1, 4), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Sum(), input_info, {}, expected1)
def test_datatype():
input_info = [([1, 2, 3, 4], "float32")]
# float
class ToFloat(Module):
def forward(self, x):
return x.float()
@tvm.script.ir_module
class expected1:
@R.function
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
(1, 2, 3, 4), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 2, 3, 4), dtype="float32") = R.astype(x, dtype="float32")
gv: R.Tensor((1, 2, 3, 4), dtype="float32") = lv
R.output(gv)
return gv
verify_model(ToFloat(), input_info, {}, expected1)
# half
class ToHalf(Module):
def forward(self, x):
return x.half()
@tvm.script.ir_module
class expected2:
@R.function
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
(1, 2, 3, 4), dtype="float16"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 2, 3, 4), dtype="float16") = R.astype(x, dtype="float16")
gv: R.Tensor((1, 2, 3, 4), dtype="float16") = lv
R.output(gv)
return gv
verify_model(ToHalf(), input_info, {}, expected2)
# type
class Type(Module):
def forward(self, x):
return x.type(torch.float32)
# type
class TypeFromAttr(Module):
def forward(self, x):
return x.type(x.getattr("dtype"))
# astype
class AsType(Module):
def forward(self, x):
return x.astype(torch.float32)
verify_model(Type(), input_info, {}, expected1)
verify_model(TypeFromAttr(), input_info, {}, expected1)
verify_model(AsType(), input_info, {}, expected1)
def test_meshgrid():
input_infos = [
(
[
3,
],
"float32",
),
(
[
3,
],
"float32",
),
]
class Meshgrid1(Module):
def forward(self, input1, input2):
return torch.meshgrid((input1, input2), indexing="ij")
class Meshgrid2(Module):
def forward(self, input1, input2):
return torch.meshgrid((input1, input2), indexing="xy")
@tvm.script.ir_module
class expected1:
@R.function
def main(
inp_0: R.Tensor((3,), dtype="float32"), inp_1: R.Tensor((3,), dtype="float32")
) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")):
with R.dataflow():
lv: R.Tuple(
R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")
) = R.meshgrid((inp_0, inp_1), indexing="ij")
gv: R.Tuple(
R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")
) = lv
R.output(gv)
return gv
@tvm.script.ir_module
class expected2:
@R.function
def main(
inp_0: R.Tensor((3,), dtype="float32"), inp_1: R.Tensor((3,), dtype="float32")
) -> R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")):
with R.dataflow():
lv: R.Tuple(
R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")
) = R.meshgrid((inp_0, inp_1), indexing="xy")
gv: R.Tuple(
R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")
) = lv
R.output(gv)
return gv
verify_model(Meshgrid1(), input_infos, {}, expected1)
verify_model(Meshgrid2(), input_infos, {}, expected2)
def test_permute():
input_info = [([1, 2, 3, 4], "float32")]
class Permute1(Module):
def forward(self, x):
return x.permute(0, 3, 2, 1)
class Permute2(Module):
def forward(self, x):
return torch.permute(x, (0, 3, 2, 1))
@tvm.script.ir_module
class expected1:
@R.function
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
(1, 4, 3, 2), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 4, 3, 2), dtype="float32") = R.permute_dims(x, axes=[0, 3, 2, 1])
gv: R.Tensor((1, 4, 3, 2), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Permute1(), input_info, {}, expected1)
verify_model(Permute2(), input_info, {}, expected1)
def test_reshape():
input_info = [([1, 2, 3, 4], "float32")]
class Reshape(Module):
def forward(self, x):
return x.reshape(2, 12)
@tvm.script.ir_module
class expected1:
@R.function
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor((2, 12), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, (2, 12))
gv: R.Tensor((2, 12), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Reshape(), input_info, {}, expected1)
def test_tile():
input_info = [([1, 3], "float32")]
class Tile1(Module):
def forward(self, x):
return x.tile((2,))
class Tile2(Module):
def forward(self, x):
return x.tile(4, 2)
class Tile3(Module):
def forward(self, x):
return torch.tile(x, (4, 2))
@tvm.script.ir_module
class expected1:
@R.function
def main(x: R.Tensor((1, 3), dtype="float32")) -> R.Tensor((1, 6), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((1, 6), dtype="float32") = R.tile(x, [2])
gv: R.Tensor((1, 6), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class expected2:
@R.function
def main(x: R.Tensor((1, 3), dtype="float32")) -> R.Tensor((4, 6), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((4, 6), dtype="float32") = R.tile(x, [4, 2])
gv: R.Tensor((4, 6), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Tile1(), input_info, {}, expected1)
verify_model(Tile2(), input_info, {}, expected2)
verify_model(Tile3(), input_info, {}, expected2)
def test_transpose():
input_info = [([1, 2, 3, 4], "float32")]
class Transpose(Module):
def forward(self, x):
return x.transpose(1, 3)
@tvm.script.ir_module
class expected1:
@R.function
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor(
(1, 4, 3, 2), dtype="float32"
):
# block 0
with R.dataflow():
lv: R.Tensor((1, 4, 3, 2), dtype="float32") = R.permute_dims(x, axes=[0, 3, 2, 1])
gv: R.Tensor((1, 4, 3, 2), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Transpose(), input_info, {}, expected1)
def test_repeat():
class Tile1(Module):
def forward(self, x: torch.Tensor):
return x.repeat(2)
class Tile2(Module):
def forward(self, x: torch.Tensor):
return x.repeat(4, 2)
@tvm.script.ir_module
class expected1:
@R.function
def main(x: R.Tensor((3,), dtype="float32")) -> R.Tensor((6,), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((6,), dtype="float32") = R.tile(x, 2)
gv: R.Tensor((6,), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class expected2:
@R.function
def main(x: R.Tensor((1, 3), dtype="float32")) -> R.Tensor((4, 6), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((4, 6), dtype="float32") = R.tile(x, [4, 2])
gv: R.Tensor((4, 6), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Tile1(), [([3], "float32")], {}, expected1)
verify_model(Tile2(), [([1, 3], "float32")], {}, expected2)
verify_model(Tile2(), [(torch.Size([1, 3]), "float32")], {}, expected2)
def test_roll():
class Roll1(Module):
def forward(self, x):
return torch.roll(x, 1)
class Roll2(Module):
def forward(self, x):
return torch.roll(x, -1, 0)
class Roll3(Module):
def forward(self, x):
return torch.roll(x, shifts=(2, 1), dims=(0, 1))
# Test case 1: torch.roll(x, 1)
@I.ir_module
class Expected1:
@R.function
def main(inp_0: R.Tensor((4, 2), dtype="int64")) -> R.Tensor((4, 2), dtype="int64"):
with R.dataflow():
lv: R.Tensor((8,), dtype="int64") = R.reshape(inp_0, R.shape([8]))
lv1: R.Tensor((7,), dtype="int64") = R.strided_slice(
lv,
axes=[0],
begin=[R.prim_value(0)],
end=[R.prim_value(7)],
strides=[R.prim_value(1)],
assume_inbound=False,
)
lv2: R.Tensor((1,), dtype="int64") = R.strided_slice(
lv,
axes=[0],
begin=[R.prim_value(7)],
end=[R.prim_value(8)],
strides=[R.prim_value(1)],
assume_inbound=False,
)
lv3: R.Tensor((8,), dtype="int64") = R.concat((lv2, lv1), axis=0)
lv4: R.Tensor((4, 2), dtype="int64") = R.reshape(lv3, R.shape([4, 2]))
gv: R.Tensor((4, 2), dtype="int64") = lv4
R.output(gv)
return gv
# Test case 2: torch.roll(x, -1, 0)
@I.ir_module
class Expected2:
@R.function
def main(inp_0: R.Tensor((4, 2), dtype="int64")) -> R.Tensor((4, 2), dtype="int64"):
with R.dataflow():
lv: R.Tensor((1, 2), dtype="int64") = R.strided_slice(
inp_0,
axes=[0],
begin=[R.prim_value(0)],
end=[R.prim_value(1)],
strides=[R.prim_value(1)],
assume_inbound=False,
)
lv1: R.Tensor((3, 2), dtype="int64") = R.strided_slice(
inp_0,
axes=[0],
begin=[R.prim_value(1)],
end=[R.prim_value(4)],
strides=[R.prim_value(1)],
assume_inbound=False,
)
lv2: R.Tensor((4, 2), dtype="int64") = R.concat((lv1, lv), axis=0)
gv: R.Tensor((4, 2), dtype="int64") = lv2
R.output(gv)
return gv
# Test case 3: torch.roll(x, shifts=(2, 1), dims=(0, 1))
@I.ir_module
class Expected3:
@R.function
def main(inp_0: R.Tensor((4, 2), dtype="int64")) -> R.Tensor((4, 2), dtype="int64"):
with R.dataflow():
lv: R.Tensor((2, 2), dtype="int64") = R.strided_slice(
inp_0,
axes=[0],
begin=[R.prim_value(0)],
end=[R.prim_value(2)],
strides=[R.prim_value(1)],
assume_inbound=False,
)
lv1: R.Tensor((2, 2), dtype="int64") = R.strided_slice(
inp_0,
axes=[0],
begin=[R.prim_value(2)],
end=[R.prim_value(4)],
strides=[R.prim_value(1)],
assume_inbound=False,
)
lv2: R.Tensor((4, 2), dtype="int64") = R.concat((lv1, lv), axis=0)
lv3: R.Tensor((4, 1), dtype="int64") = R.strided_slice(
lv2,
axes=[1],
begin=[R.prim_value(0)],
end=[R.prim_value(1)],
strides=[R.prim_value(1)],
assume_inbound=False,
)
lv4: R.Tensor((4, 1), dtype="int64") = R.strided_slice(
lv2,
axes=[1],
begin=[R.prim_value(1)],
end=[R.prim_value(2)],
strides=[R.prim_value(1)],
assume_inbound=False,
)
lv5: R.Tensor((4, 2), dtype="int64") = R.concat((lv4, lv3), axis=1)
gv: R.Tensor((4, 2), dtype="int64") = lv5
R.output(gv)
return gv
input_info = [([4, 2], "int64")]
verify_model(Roll1(), input_info, {}, Expected1)
verify_model(Roll2(), input_info, {}, Expected2)
verify_model(Roll3(), input_info, {}, Expected3)
def test_view():
input_info = [([1, 2, 3, 4], "float32")]
class View(Module):
def forward(self, x):
return x.view(2, 12)
@tvm.script.ir_module
class expected1:
@R.function
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tensor((2, 12), dtype="float32"):
# block 0
with R.dataflow():
lv: R.Tensor((2, 12), dtype="float32") = R.reshape(x, (2, 12))
gv: R.Tensor((2, 12), dtype="float32") = lv
R.output(gv)
return gv
verify_model(View(), input_info, {}, expected1)
def test_keep_params():
class Conv2D1(Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(3, 6, 7, bias=True)
def forward(self, input):
return self.conv(input)
@tvm.script.ir_module
class expected1:
@R.function
def main(
input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
w1: R.Tensor((6,), dtype="float32"),
w2: R.Tensor((6, 3, 7, 7), dtype="float32"),
) -> R.Tensor((1, 6, 4, 4), dtype="float32"):
R.func_attr({"num_input": 1})
# block 0
with R.dataflow():
lv1: R.Tensor((1, 6, 4, 4), dtype="float32") = R.nn.conv2d(
input_1,
w2,
strides=[1, 1],
padding=[0, 0, 0, 0],
dilation=[1, 1],
data_layout="NCHW",
kernel_layout="OIHW",
out_layout="NCHW",
out_dtype="float32",
)
lv2: R.Tensor((1, 6, 1, 1), dtype="float32") = R.reshape(w1, [1, 6, 1, 1])
lv3: R.Tensor((1, 6, 4, 4), dtype="float32") = R.add(lv1, lv2)
gv: R.Tensor((1, 6, 4, 4), dtype="float32") = lv3
R.output(gv)
return gv
model = Conv2D1()
graph_model = fx.symbolic_trace(model)
mod = from_fx(graph_model, [([1, 3, 10, 10], "float32")], keep_params_as_input=True)
mod, params = detach_params(mod)
tvm.ir.assert_structural_equal(mod, expected1)
func = mod["main"]
params = params["main"]
assert len(params) == len(func.params) - 1
for param_var, param_tensor in zip(func.params[1:], params):
assert tuple(x.value for x in param_var.ty.shape.values) == param_tensor.shape
assert param_var.ty.dtype == param_tensor.dtype
tvm.testing.assert_allclose(params[0].numpy(), model.conv.bias.detach().detach().numpy())
tvm.testing.assert_allclose(params[1].numpy(), model.conv.weight.detach().detach().numpy())
def test_unwrap_unit_return_tuple():
class Identity(Module):
def __init__(self):
super().__init__()
def forward(self, x):
return (x,)
@tvm.script.ir_module
class Expected:
@R.function
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor(
(256, 256), dtype="float32"
):
with R.dataflow():
gv: R.Tensor((256, 256), dtype="float32") = inp_0
R.output(gv)
return gv
graph_model = fx.symbolic_trace(Identity())
mod = from_fx(graph_model, [([256, 256], "float32")], unwrap_unit_return_tuple=True)
tvm.ir.assert_structural_equal(mod, Expected)
def test_no_bind_return_tuple():
class Identity(Module):
def __init__(self):
super().__init__()
def forward(self, x, y):
return (x, y)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((256, 256), dtype="float32"),
inp_1: R.Tensor((256, 256), dtype="float32"),
) -> R.Tuple(R.Tensor((256, 256), dtype="float32"), R.Tensor((256, 256), dtype="float32")):
with R.dataflow():
gv: R.Tensor((256, 256), dtype="float32") = inp_0
gv1: R.Tensor((256, 256), dtype="float32") = inp_1
R.output(gv, gv1)
return (gv, gv1)
graph_model = fx.symbolic_trace(Identity())
mod = from_fx(
graph_model, [([256, 256], "float32"), ([256, 256], "float32")], no_bind_return_tuple=True
)
tvm.ir.assert_structural_equal(mod, Expected)
def test_argmax():
class Argmax1(Module):
def __init__(self) -> None:
super().__init__()
def forward(self, input):
return torch.argmax(input, dim=-1)
class Argmax2(Module):
def __init__(self) -> None:
super().__init__()
def forward(self, input):
return torch.argmax(input, dim=-1, keepdim=True)
@tvm.script.ir_module
class Expected1:
@R.function
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((256,), dtype="int64"):
with R.dataflow():
lv: R.Tensor((256,), dtype="int64") = R.argmax(inp_0, axis=-1, keepdims=False)
gv: R.Tensor((256,), dtype="int64") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class Expected2:
@R.function
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((256, 1), dtype="int64"):
with R.dataflow():
lv: R.Tensor((256, 1), dtype="int64") = R.argmax(inp_0, axis=-1, keepdims=True)
gv: R.Tensor((256, 1), dtype="int64") = lv
R.output(gv)
return gv
verify_model(Argmax1(), [([256, 256], "float32")], {}, Expected1)
verify_model(Argmax2(), [([256, 256], "float32")], {}, Expected2)
def test_argmin():
class Argmin1(Module):
def __init__(self) -> None:
super().__init__()
def forward(self, input):
return torch.argmin(input)
class Argmin2(Module):
def __init__(self) -> None:
super().__init__()
def forward(self, input):
return torch.argmin(input, keepdim=True)
@tvm.script.ir_module
class Expected1:
@R.function
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((), dtype="int64"):
with R.dataflow():
lv: R.Tensor((), dtype="int64") = R.argmin(inp_0, axis=None, keepdims=False)
gv: R.Tensor((), dtype="int64") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class Expected2:
@R.function
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((1, 1), dtype="int64"):
with R.dataflow():
lv: R.Tensor((1, 1), dtype="int64") = R.argmin(inp_0, axis=None, keepdims=True)
gv: R.Tensor((1, 1), dtype="int64") = lv
R.output(gv)
return gv
verify_model(Argmin1(), [([256, 256], "float32")], {}, Expected1)
verify_model(Argmin2(), [([256, 256], "float32")], {}, Expected2)
def test_to():
class To1(Module):
def forward(self, input):
return input.to(torch.float16)
class To2(Module):
def forward(self, input):
return input.to("cpu")
@I.ir_module
class Expected1:
@R.function
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor(
(256, 256), dtype="float16"
):
with R.dataflow():
lv: R.Tensor((256, 256), dtype="float16") = R.astype(inp_0, dtype="float16")
gv: R.Tensor((256, 256), dtype="float16") = lv
R.output(gv)
return gv
@I.ir_module
class Expected2:
@R.function
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor(
(256, 256), dtype="float32"
):
with R.dataflow():
gv: R.Tensor((256, 256), dtype="float32") = inp_0
R.output(gv)
return gv
verify_model(To1(), [([256, 256], "float32")], {}, Expected1)
verify_model(To2(), [([256, 256], "float32")], {}, Expected2)
def test_mean():
class Mean(Module):
def forward(self, input):
return input.mean(-1)
class MeanKeepDim(Module):
def forward(self, input):
return input.mean(-1, keepdim=True)
@I.ir_module
class Expected1:
@R.function
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor((256,), dtype="float32"):
with R.dataflow():
lv: R.Tensor((256,), dtype="float32") = R.mean(inp_0, axis=[-1], keepdims=False)
gv: R.Tensor((256,), dtype="float32") = lv
R.output(gv)
return gv
@I.ir_module
class Expected2:
@R.function
def main(inp_0: R.Tensor((256, 256), dtype="float32")) -> R.Tensor(
(256, 1), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((256, 1), dtype="float32") = R.mean(inp_0, axis=[-1], keepdims=True)
gv: R.Tensor((256, 1), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Mean(), [([256, 256], "float32")], {}, Expected1)
verify_model(MeanKeepDim(), [([256, 256], "float32")], {}, Expected2)
def test_cat():
class Cat0(Module):
def forward(self, x, y):
return torch.cat((x, y))
class Cat1(Module):
def forward(self, x, y):
return torch.cat((x, y), dim=1)
class Cat2(Module):
def forward(self, x, y):
return torch.cat((x, y), 1)
class Cat3(Module):
def forward(self, x, y):
return torch.concat((x, y), dim=0)
@I.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((2, 3), dtype="float32"),
inp_1: R.Tensor((2, 3), dtype="float32"),
) -> R.Tensor((4, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((4, 3), dtype="float32") = R.concat((inp_0, inp_1), axis=0)
gv: R.Tensor((4, 3), dtype="float32") = lv
R.output(gv)
return gv
@I.ir_module
class Expected2:
@R.function
def main(
inp_0: R.Tensor((2, 3), dtype="float32"),
inp_1: R.Tensor((2, 3), dtype="float32"),
) -> R.Tensor((2, 6), dtype="float32"):
with R.dataflow():
lv: R.Tensor((2, 6), dtype="float32") = R.concat((inp_0, inp_1), axis=1)
gv: R.Tensor((2, 6), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Cat0(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected1)
verify_model(Cat1(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected2)
verify_model(Cat2(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected2)
verify_model(Cat3(), [([2, 3], "float32"), ([2, 3], "float32")], {}, Expected1)
def test_max():
class Max(Module):
def forward(self, x, y):
return torch.max(x, y)
@I.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((256, 256), dtype="float32"),
inp_1: R.Tensor((256, 256), dtype="float32"),
) -> R.Tensor((256, 256), dtype="float32"):
with R.dataflow():
lv: R.Tensor((256, 256), dtype="float32") = R.maximum(inp_0, inp_1)
gv: R.Tensor((256, 256), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Max(), [([256, 256], "float32"), ([256, 256], "float32")], {}, Expected1)
def test_min():
class Min(Module):
def forward(self, x, y):
return torch.min(x, y)
@I.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((256, 256), dtype="float32"),
inp_1: R.Tensor((256, 256), dtype="float32"),
) -> R.Tensor((256, 256), dtype="float32"):
with R.dataflow():
lv: R.Tensor((256, 256), dtype="float32") = R.minimum(inp_0, inp_1)
gv: R.Tensor((256, 256), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Min(), [([256, 256], "float32"), ([256, 256], "float32")], {}, Expected1)
def test_atan2():
class Atan2(Module):
def forward(self, x, y):
return torch.atan2(x, y)
@I.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((256, 256), dtype="float32"),
inp_1: R.Tensor((256, 256), dtype="float32"),
) -> R.Tensor((256, 256), dtype="float32"):
with R.dataflow():
lv: R.Tensor((256, 256), dtype="float32") = R.atan2(inp_0, inp_1)
gv: R.Tensor((256, 256), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Atan2(), [([256, 256], "float32"), ([256, 256], "float32")], {}, Expected1)
def test_attention():
@I.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((32, 8, 128, 64), dtype="float32"),
inp_1: R.Tensor((32, 8, 128, 64), dtype="float32"),
inp_2: R.Tensor((32, 8, 128, 64), dtype="float32"),
) -> R.Tensor((32, 8, 128, 64), dtype="float32"):
with R.dataflow():
lv: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
inp_0, axes=[0, 2, 1, 3]
)
lv1: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
inp_1, axes=[0, 2, 1, 3]
)
lv2: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
inp_2, axes=[0, 2, 1, 3]
)
lv3: R.Tensor((32, 128, 8, 64), dtype="float32") = R.nn.attention(
lv, lv1, lv2, scale=None
)
lv4: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
lv3, axes=[0, 2, 1, 3]
)
gv: R.Tensor((32, 8, 128, 64), dtype="float32") = lv4
R.output(gv)
return gv
@I.ir_module
class Expected2:
@R.function
def main(
inp_0: R.Tensor((32, 8, 128, 64), dtype="float32"),
inp_1: R.Tensor((32, 8, 128, 64), dtype="float32"),
inp_2: R.Tensor((32, 8, 128, 64), dtype="float32"),
inp_3: R.Tensor((32, 8, 128, 128), dtype="float32"),
) -> R.Tensor((32, 8, 128, 64), dtype="float32"):
with R.dataflow():
lv: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
inp_0, axes=[0, 2, 1, 3]
)
lv1: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
inp_1, axes=[0, 2, 1, 3]
)
lv2: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
inp_2, axes=[0, 2, 1, 3]
)
lv3: R.Tensor((32, 128, 8, 64), dtype="float32") = R.nn.attention(
lv, lv1, lv2, inp_3, scale=None
)
lv4: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
lv3, axes=[0, 2, 1, 3]
)
gv: R.Tensor((32, 8, 128, 64), dtype="float32") = lv4
R.output(gv)
return gv
@I.ir_module
class Expected3:
@R.function
def main(
inp_0: R.Tensor((32, 8, 128, 64), dtype="float32"),
inp_1: R.Tensor((32, 8, 128, 64), dtype="float32"),
inp_2: R.Tensor((32, 8, 128, 64), dtype="float32"),
) -> R.Tensor((32, 8, 128, 64), dtype="float32"):
with R.dataflow():
lv: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
inp_0, axes=[0, 2, 1, 3]
)
lv1: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
inp_1, axes=[0, 2, 1, 3]
)
lv2: R.Tensor((32, 128, 8, 64), dtype="float32") = R.permute_dims(
inp_2, axes=[0, 2, 1, 3]
)
lv3: R.Tensor((32, 128, 8, 64), dtype="float32") = R.nn.attention(
lv, lv1, lv2, scale=None, causal_mask="TopLeft"
)
lv4: R.Tensor((32, 8, 128, 64), dtype="float32") = R.permute_dims(
lv3, axes=[0, 2, 1, 3]
)
gv: R.Tensor((32, 8, 128, 64), dtype="float32") = lv4
R.output(gv)
return gv
verify_model(
lambda q, k, v: F.scaled_dot_product_attention(q, k, v),
[
([32, 8, 128, 64], "float32"),
([32, 8, 128, 64], "float32"),
([32, 8, 128, 64], "float32"),
],
{},
Expected1,
)
verify_model(
lambda q, k, v, mask: F.scaled_dot_product_attention(q, k, v, mask),
[
([32, 8, 128, 64], "float32"),
([32, 8, 128, 64], "float32"),
([32, 8, 128, 64], "float32"),
([32, 8, 128, 128], "float32"),
],
{},
Expected2,
)
verify_model(
lambda q, k, v: F.scaled_dot_product_attention(q, k, v, is_causal=True),
[
([32, 8, 128, 64], "float32"),
([32, 8, 128, 64], "float32"),
([32, 8, 128, 64], "float32"),
],
{},
Expected3,
)
def test_sym_size_int():
class SymSizeInt1(Module):
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, x):
return torch.ops.aten.sym_size.int(x, self.dim)
@I.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((1, 3, 4), dtype="float32"),
) -> R.Tensor((), dtype="int32"):
with R.dataflow():
lv: R.Tensor((), dtype="int32") = R.const(3, "int32")
gv: R.Tensor((), dtype="int32") = lv
R.output(gv)
return gv
verify_model(SymSizeInt1(dim=1), [([1, 3, 4], "float32")], {}, Expected1)
verify_model(SymSizeInt1(dim=-2), [([1, 3, 4], "float32")], {}, Expected1)
def test_stack():
input_info = [
([1, 3, 10, 10], "float32"),
([1, 3, 10, 10], "float32"),
([1, 3, 10, 10], "float32"),
]
class Stack(Module):
def forward(self, data, data1, data2):
return torch.stack((data, data1, data2), dim=0)
@tvm.script.ir_module
class expected:
@R.function
def main(
inp_0: R.Tensor((1, 3, 10, 10), dtype="float32"),
inp_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
inp_2: R.Tensor((1, 3, 10, 10), dtype="float32"),
) -> R.Tensor((3, 1, 3, 10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((3, 1, 3, 10, 10), dtype="float32") = R.stack(
(inp_0, inp_1, inp_2), axis=0
)
gv: R.Tensor((3, 1, 3, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Stack(), input_info, {}, expected)
def test_scatter():
input_info = [([20, 20], "float32"), ([2, 5], "int64"), ([2, 5], "float32")]
class Scatter(Module):
def forward(self, data, index, src):
return data.scatter(dim=0, index=index, src=src)
@tvm.script.ir_module
class expected:
@R.function
def main(
inp_0: R.Tensor((20, 20), dtype="float32"),
inp_1: R.Tensor((2, 5), dtype="int64"),
inp_2: R.Tensor((2, 5), dtype="float32"),
) -> R.Tensor((20, 20), dtype="float32"):
with R.dataflow():
lv: R.Tensor((20, 20), dtype="float32") = R.scatter_elements(
inp_0, inp_1, inp_2, axis=0, reduction="update"
)
gv: R.Tensor((20, 20), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Scatter(), input_info, {}, expected)
def test_slice_scatter():
class SliceScatter1(Module):
def forward(self, input, src):
return torch.slice_scatter(input, src, dim=1, start=1, end=7, step=2)
@tvm.script.ir_module
class expected1:
@R.function
def main(
a: R.Tensor((8, 8, 10, 10), dtype="float32"),
b: R.Tensor((8, 3, 10, 10), dtype="float32"),
) -> R.Tensor((8, 8, 10, 10), dtype="float32"):
with R.dataflow():
lv: R.Tensor((8, 8, 10, 10), dtype="float32") = R.slice_scatter(
a, b, R.prim_value(1), R.prim_value(7), R.prim_value(2), axis=1
)
gv: R.Tensor((8, 8, 10, 10), dtype="float32") = lv
R.output(gv)
return gv
class SliceScatter2(Module):
def forward(self, input, src):
return torch.slice_scatter(input, src, dim=0, start=0, end=6, step=1)
@I.ir_module
class expected2:
@R.function
def main(
a: R.Tensor((8, 16), dtype="float32"), b: R.Tensor((6, 16), dtype="float32")
) -> R.Tensor((8, 16), dtype="float32"):
with R.dataflow():
lv: R.Tensor((8, 16), dtype="float32") = R.slice_scatter(
a, b, R.prim_value(0), R.prim_value(6), R.prim_value(1), axis=0
)
gv: R.Tensor((8, 16), dtype="float32") = lv
R.output(gv)
return gv
class SliceScatterNegative(Module):
def forward(self, input, src):
return torch.slice_scatter(input, src, dim=1, start=0, end=-2, step=1)
@tvm.script.ir_module
class expected_slice_scatter:
@R.function
def main(
a: R.Tensor((2, 5), dtype="float32"), b: R.Tensor((2, 3), dtype="float32")
) -> R.Tensor((2, 5), dtype="float32"):
with R.dataflow():
lv: R.Tensor((2, 5), dtype="float32") = R.slice_scatter(
a, b, R.prim_value(0), R.prim_value(3), R.prim_value(1), axis=1
)
gv: R.Tensor((2, 5), dtype="float32") = lv
R.output(gv)
return gv
verify_model(
SliceScatter1(), [((8, 8, 10, 10), "float32"), ((8, 3, 10, 10), "float32")], {}, expected1
)
verify_model(SliceScatter2(), [((8, 16), "float32"), ((6, 16), "float32")], {}, expected2)
verify_model(
SliceScatterNegative(),
[((2, 5), "float32"), ((2, 3), "float32")],
{},
expected_slice_scatter,
)
def test_masked_scatter():
class MaskedScatter1(Module):
def forward(self, data, mask, src):
return data.masked_scatter(mask, src)
class MaskedScatter2(Module):
def forward(self, data, mask, src):
return data.masked_scatter(mask, src)
@tvm.script.ir_module
class expected1:
@R.function
def main(
inp_0: R.Tensor((5,), dtype="float32"),
inp_1: R.Tensor((5,), dtype="bool"),
inp_2: R.Tensor((10,), dtype="float32"),
) -> R.Tensor((5,), dtype="float32"):
with R.dataflow():
lv: R.Tensor((5,), dtype="int32") = R.cumsum(
inp_1, axis=0, dtype="int32", exclusive=False
)
lv1: R.Tensor((5,), dtype="int32") = R.subtract(lv, R.const(1, "int32"))
lv2: R.Tensor((5,), dtype="float32") = R.take(inp_2, lv1, axis=0)
lv3: R.Tensor((5,), dtype="float32") = R.where(inp_1, lv2, inp_0)
gv: R.Tensor((5,), dtype="float32") = lv3
R.output(gv)
return gv
@tvm.script.ir_module
class expected2:
@R.function
def main(
inp_0: R.Tensor((2, 5), dtype="float32"),
inp_1: R.Tensor((2, 5), dtype="bool"),
inp_2: R.Tensor((3, 5), dtype="float32"),
) -> R.Tensor((2, 5), dtype="float32"):
with R.dataflow():
lv: R.Tensor((10,), dtype="bool") = R.reshape(inp_1, R.shape([10]))
lv1: R.Tensor((10,), dtype="int32") = R.cumsum(
lv, axis=0, dtype="int32", exclusive=False
)
lv2: R.Tensor((10,), dtype="int32") = R.subtract(lv1, R.const(1, "int32"))
lv3: R.Tensor((15,), dtype="float32") = R.reshape(inp_2, R.shape([15]))
lv4: R.Tensor((10,), dtype="float32") = R.take(lv3, lv2, axis=0)
lv5: R.Tensor((2, 5), dtype="float32") = R.reshape(lv4, R.shape([2, 5]))
lv6: R.Tensor((2, 5), dtype="float32") = R.where(inp_1, lv5, inp_0)
gv: R.Tensor((2, 5), dtype="float32") = lv6
R.output(gv)
return gv
verify_model(
MaskedScatter1(), [([5], "float32"), ([5], "bool"), ([10], "float32")], {}, expected1
)
verify_model(
MaskedScatter2(),
[([2, 5], "float32"), ([2, 5], "bool"), ([3, 5], "float32")],
{},
expected2,
)
def test_is_floating_point():
class IsFloatingPoint(Module):
def forward(self, x):
return torch.is_floating_point(x)
@tvm.script.ir_module
class Expected:
@R.function
def main(inp_0: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((), dtype="bool"):
with R.dataflow():
gv: R.Tensor((), dtype="bool") = R.const(True, "bool")
R.output(gv)
return gv
verify_model(IsFloatingPoint(), [([2, 3], "float32")], {}, Expected)
def test_gather():
class Gather0(Module):
def forward(self, data, indices):
return torch.gather(data, 0, indices)
class Gather1(Module):
def forward(self, data, indices):
return torch.gather(data, 1, indices)
class Gather2(Module):
def forward(self, data, indices):
return torch.gather(data, -1, indices)
class Gather3(Module):
def forward(self, data, indices):
return torch.gather(data, -2, indices)
@tvm.script.ir_module
class Expected0:
@R.function
def main(
inp_0: R.Tensor((2, 3), dtype="float32"),
inp_1: R.Tensor((2, 3), dtype="int32"),
) -> R.Tensor((2, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=0)
gv: R.Tensor((2, 3), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((2, 3), dtype="float32"),
inp_1: R.Tensor((2, 3), dtype="int32"),
) -> R.Tensor((2, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=1)
gv: R.Tensor((2, 3), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class Expected2:
@R.function
def main(
inp_0: R.Tensor((2, 3), dtype="float32"),
inp_1: R.Tensor((2, 3), dtype="int32"),
) -> R.Tensor((2, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=-1)
gv: R.Tensor((2, 3), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class Expected3:
@R.function
def main(
inp_0: R.Tensor((2, 3), dtype="float32"),
inp_1: R.Tensor((2, 3), dtype="int32"),
) -> R.Tensor((2, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((2, 3), dtype="float32") = R.gather_elements(inp_0, inp_1, axis=-2)
gv: R.Tensor((2, 3), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Gather0(), [([2, 3], "float32"), ([2, 3], "int32")], {}, Expected0)
verify_model(Gather1(), [([2, 3], "float32"), ([2, 3], "int32")], {}, Expected1)
verify_model(Gather2(), [([2, 3], "float32"), ([2, 3], "int32")], {}, Expected2)
verify_model(Gather3(), [([2, 3], "float32"), ([2, 3], "int32")], {}, Expected3)
def test_index_put():
# Test case 1: 1D input
class IndexPut1D(Module):
def forward(self, data, indices_0, values):
indices_tuple = (indices_0,)
return data.index_put_(indices_tuple, values, accumulate=False)
input_info_1d = [((64,), "float32"), ((128,), "int64"), ((128,), "float32")]
@I.ir_module
class Expected1D:
@R.function
def main(
data: R.Tensor((64,), dtype="float32"),
indices_0: R.Tensor((128,), dtype="int64"),
values: R.Tensor((128,), dtype="float32"),
) -> R.Tensor((64,), dtype="float32"):
with R.dataflow():
lv: R.Tensor((64,), dtype="float32") = R.index_put(
data, R.tuple(indices_0), values, accumulate=False
)
gv: R.Tensor((64,), dtype="float32") = lv
R.output(gv)
return gv
# Test case 2: 2D input
class IndexPut2D(Module):
def forward(self, data, indices_0, indices_1, values):
indices_tuple = (indices_0, indices_1)
return data.index_put_(indices_tuple, values, accumulate=False)
input_info_2d = [
((32, 64), "float32"),
((128,), "int64"),
((128,), "int64"),
((128,), "float32"),
]
@I.ir_module
class Expected2D:
@R.function
def main(
data: R.Tensor((32, 64), dtype="float32"),
indices_0: R.Tensor((128,), dtype="int64"),
indices_1: R.Tensor((128,), dtype="int64"),
values: R.Tensor((128,), dtype="float32"),
) -> R.Tensor((32, 64), dtype="float32"):
with R.dataflow():
lv: R.Tensor((32, 64), dtype="float32") = R.index_put(
data, R.tuple(indices_0, indices_1), values, accumulate=False
)
gv: R.Tensor((32, 64), dtype="float32") = lv
R.output(gv)
return gv
# Test case 3: 3D input
class IndexPut3D(Module):
def forward(self, data, indices_0, indices_1, indices_2, values):
indices_tuple = (indices_0, indices_1, indices_2)
return data.index_put_(indices_tuple, values, accumulate=False)
input_info_3d = [
((16, 32, 64), "float32"),
((128,), "int64"),
((128,), "int64"),
((128,), "int64"),
((128,), "float32"),
]
@I.ir_module
class Expected3D:
@R.function
def main(
data: R.Tensor((16, 32, 64), dtype="float32"),
indices_0: R.Tensor((128,), dtype="int64"),
indices_1: R.Tensor((128,), dtype="int64"),
indices_2: R.Tensor((128,), dtype="int64"),
values: R.Tensor((128,), dtype="float32"),
) -> R.Tensor((16, 32, 64), dtype="float32"):
with R.dataflow():
lv: R.Tensor((16, 32, 64), dtype="float32") = R.index_put(
data, R.tuple(indices_0, indices_1, indices_2), values, accumulate=False
)
gv: R.Tensor((16, 32, 64), dtype="float32") = lv
R.output(gv)
return gv
# Test case 4: 4D input
class IndexPut4D(Module):
def forward(self, data, indices_0, indices_1, indices_2, indices_3, values):
indices_tuple = (indices_0, indices_1, indices_2, indices_3)
return data.index_put_(indices_tuple, values, accumulate=False)
input_info_4d = [
((8, 16, 32, 64), "float32"),
((128,), "int64"),
((128,), "int64"),
((128,), "int64"),
((128,), "int64"),
((128,), "float32"),
]
@I.ir_module
class Expected4D:
@R.function
def main(
data: R.Tensor((8, 16, 32, 64), dtype="float32"),
indices_0: R.Tensor((128,), dtype="int64"),
indices_1: R.Tensor((128,), dtype="int64"),
indices_2: R.Tensor((128,), dtype="int64"),
indices_3: R.Tensor((128,), dtype="int64"),
values: R.Tensor((128,), dtype="float32"),
) -> R.Tensor((8, 16, 32, 64), dtype="float32"):
with R.dataflow():
lv: R.Tensor((8, 16, 32, 64), dtype="float32") = R.index_put(
data,
R.tuple(indices_0, indices_1, indices_2, indices_3),
values,
accumulate=False,
)
gv: R.Tensor((8, 16, 32, 64), dtype="float32") = lv
R.output(gv)
return gv
# Test case 5: 5D input
class IndexPut5D(Module):
def forward(self, data, indices_0, indices_1, indices_2, indices_3, indices_4, values):
indices_tuple = (indices_0, indices_1, indices_2, indices_3, indices_4)
return data.index_put_(indices_tuple, values, accumulate=False)
input_info_5d = [
((4, 8, 16, 32, 64), "float32"),
((128,), "int64"),
((128,), "int64"),
((128,), "int64"),
((128,), "int64"),
((128,), "int64"),
((128,), "float32"),
]
@I.ir_module
class Expected5D:
@R.function
def main(
data: R.Tensor((4, 8, 16, 32, 64), dtype="float32"),
indices_0: R.Tensor((128,), dtype="int64"),
indices_1: R.Tensor((128,), dtype="int64"),
indices_2: R.Tensor((128,), dtype="int64"),
indices_3: R.Tensor((128,), dtype="int64"),
indices_4: R.Tensor((128,), dtype="int64"),
values: R.Tensor((128,), dtype="float32"),
) -> R.Tensor((4, 8, 16, 32, 64), dtype="float32"):
with R.dataflow():
lv: R.Tensor((4, 8, 16, 32, 64), dtype="float32") = R.index_put(
data,
R.tuple(indices_0, indices_1, indices_2, indices_3, indices_4),
values,
accumulate=False,
)
gv: R.Tensor((4, 8, 16, 32, 64), dtype="float32") = lv
R.output(gv)
return gv
# Run verification for each case
verify_model(IndexPut1D(), input_info_1d, {}, Expected1D)
verify_model(IndexPut2D(), input_info_2d, {}, Expected2D)
verify_model(IndexPut3D(), input_info_3d, {}, Expected3D)
verify_model(IndexPut4D(), input_info_4d, {}, Expected4D)
verify_model(IndexPut5D(), input_info_5d, {}, Expected5D)
def test_flip():
class Flip0(Module):
def forward(self, data):
return torch.flip(data, [0])
class Flip1(Module):
def forward(self, data):
return torch.flip(data, [1])
@tvm.script.ir_module
class Expected0:
@R.function
def main(
inp_0: R.Tensor((2, 2), dtype="float32"),
) -> R.Tensor((2, 2), dtype="float32"):
with R.dataflow():
lv: R.Tensor((2, 2), dtype="float32") = R.flip(inp_0, axis=0)
gv: R.Tensor((2, 2), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((2, 2), dtype="float32"),
) -> R.Tensor((2, 2), dtype="float32"):
with R.dataflow():
lv: R.Tensor((2, 2), dtype="float32") = R.flip(inp_0, axis=1)
gv: R.Tensor((2, 2), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Flip0(), [([2, 2], "float32")], {}, Expected0)
verify_model(Flip1(), [([2, 2], "float32")], {}, Expected1)
def test_flip_multi_axis():
class FlipMulti(Module):
def forward(self, data):
return torch.flip(data, [0, 1])
@tvm.script.ir_module
class ExpectedMulti:
@R.function
def main(
inp_0: R.Tensor((2, 3), dtype="float32"),
) -> R.Tensor((2, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((2, 3), dtype="float32") = R.flip(inp_0, axis=0)
lv1: R.Tensor((2, 3), dtype="float32") = R.flip(lv, axis=1)
gv: R.Tensor((2, 3), dtype="float32") = lv1
R.output(gv)
return gv
verify_model(FlipMulti(), [([2, 3], "float32")], {}, ExpectedMulti)
def test_take():
class Take(Module):
def forward(self, data, indices):
return torch.take(data, indices)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5,), dtype="float32"),
inp_1: R.Tensor((3,), dtype="int32"),
) -> R.Tensor((3,), dtype="float32"):
with R.dataflow():
lv: R.Tensor((3,), dtype="int32") = R.astype(inp_1, "int32")
lv1: R.Tensor((3,), dtype="float32") = R.take(inp_0, lv)
gv: R.Tensor((3,), dtype="float32") = lv1
R.output(gv)
return gv
verify_model(Take(), [([5], "float32"), ([3], "int32")], {}, Expected)
def test_one_hot():
class OneHot(Module):
def forward(self, indices):
return torch.nn.functional.one_hot(indices, num_classes=10)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5,), dtype="int32"),
) -> R.Tensor((5, 10), dtype="int64"):
with R.dataflow():
lv: R.Tensor((5, 10), dtype="int64") = R.one_hot(
inp_0, R.prim_value(1), R.prim_value(0), depth=10, axis=-1
)
gv: R.Tensor((5, 10), dtype="int64") = lv
R.output(gv)
return gv
verify_model(OneHot(), [([5], "int32")], {}, Expected)
def test_empty_like():
class EmptyLike(Module):
def forward(self, data):
return torch.empty_like(data)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5,), dtype="float32"),
) -> R.Tensor((5,), dtype="float32"):
with R.dataflow():
lv: R.Tensor((5,), dtype="float32") = R.zeros_like(inp_0)
gv: R.Tensor((5,), dtype="float32") = lv
R.output(gv)
return gv
verify_model(EmptyLike(), [([5], "float32")], {}, Expected)
def test_ones_like():
class OnesLike(Module):
def forward(self, data):
return torch.ones_like(data)
@tvm.script.ir_module
class Expected:
@R.function
def main(inp_0: R.Tensor((128, 128), dtype="float32")) -> R.Tensor(
(128, 128), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((128, 128), dtype="float32") = R.ones_like(inp_0)
gv: R.Tensor((128, 128), dtype="float32") = lv
R.output(gv)
return gv
verify_model(OnesLike(), [([128, 128], "float32")], {}, Expected)
def test_zero_inplace():
class ZeroInplace(Module):
def forward(self, data):
return data.zero_()
@tvm.script.ir_module
class Expected:
@R.function
def main(inp_0: R.Tensor((128, 128), dtype="float32")) -> R.Tensor(
(128, 128), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((128, 128), dtype="float32") = R.zeros_like(inp_0)
gv: R.Tensor((128, 128), dtype="float32") = lv
R.output(gv)
return gv
verify_model(ZeroInplace(), [([128, 128], "float32")], {}, Expected)
def test_zeros_like():
class ZerosLike(Module):
def forward(self, data):
return torch.zeros_like(data)
@tvm.script.ir_module
class Expected:
@R.function
def main(inp_0: R.Tensor((128, 128), dtype="float32")) -> R.Tensor(
(128, 128), dtype="float32"
):
with R.dataflow():
lv: R.Tensor((128, 128), dtype="float32") = R.zeros_like(inp_0)
gv: R.Tensor((128, 128), dtype="float32") = lv
R.output(gv)
return gv
verify_model(ZerosLike(), [([128, 128], "float32")], {}, Expected)
def test_type_as():
class TypeAs(Module):
def forward(self, data, other):
return data.type_as(other)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((128, 128), dtype="float16"),
inp_1: R.Tensor((128, 128), dtype="float32"),
) -> R.Tensor((128, 128), dtype="float32"):
with R.dataflow():
lv: R.Tensor((128, 128), dtype="float32") = R.astype(inp_0, dtype="float32")
gv: R.Tensor((128, 128), dtype="float32") = lv
R.output(gv)
return gv
verify_model(TypeAs(), [([128, 128], "float16"), ([128, 128], "float32")], {}, Expected)
def test_item():
class Item(Module):
def forward(self, data):
return data.item()
@tvm.script.ir_module
class Expected:
@R.function
def main(inp_0: R.Tensor((1,), dtype="float32")) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((), dtype="float32") = R.take(inp_0, R.const(0, "int64"), axis=0)
gv: R.Tensor((), dtype="float32") = lv
R.output(gv)
return gv
verify_model(
Item(),
[
(
[1],
"float32",
)
],
{},
Expected,
)
def test_numel():
class Numel(Module):
def forward(self, data):
return torch.numel(data)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="float32"),
) -> R.Tensor((), dtype="int32"):
with R.dataflow():
gv: R.Tensor((), dtype="int32") = R.const(15, "int32")
R.output(gv)
return gv
verify_model(Numel(), [([5, 3], "float32")], {}, Expected)
def test_select():
class Select(Module):
def forward(self, data):
return torch.select(data, 0, 1)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="float32"),
) -> R.Tensor((3,), dtype="float32"):
with R.dataflow():
lv: R.Tensor((3,), dtype="float32") = R.take(inp_0, R.const(1, "int64"), axis=0)
gv: R.Tensor((3,), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Select(), [([5, 3], "float32")], {}, Expected)
def test_inplace_copy():
class Inplace_Copy(Module):
def forward(self, x, y):
x.copy_(y)
return x
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((1, 2, 3, 4), dtype="float32"),
inp_1: R.Tensor((1, 2, 3, 4), dtype="float32"),
) -> R.Tensor((1, 2, 3, 4), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 2, 3, 4), dtype="float32") = R.broadcast_to(
inp_1, R.shape([1, 2, 3, 4])
)
gv: R.Tensor((1, 2, 3, 4), dtype="float32") = lv
R.output(gv)
return gv
class CopyBroadcast(Module):
def forward(self, x, src):
x.copy_(src)
return x
@tvm.script.ir_module
class expected_copy:
@R.function
def main(
x: R.Tensor((2, 3), dtype="float32"), src: R.Tensor((), dtype="int64")
) -> R.Tensor((2, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((), dtype="float32") = R.astype(src, dtype="float32")
lv1: R.Tensor((2, 3), dtype="float32") = R.broadcast_to(lv, (2, 3))
gv: R.Tensor((2, 3), dtype="float32") = lv1
R.output(gv)
return gv
verify_model(
Inplace_Copy(),
[((1, 2, 3, 4), "float32"), ((1, 2, 3, 4), "float32")],
{},
Expected,
)
verify_model(CopyBroadcast(), [((2, 3), "float32"), ((), "int64")], {}, expected_copy)
def test_clone():
class Clone(Module):
def forward(self, x):
return x.clone()
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="float32"),
) -> R.Tensor((5, 3), dtype="float32"):
with R.dataflow():
gv: R.Tensor((5, 3), dtype="float32") = inp_0
R.output(gv)
return gv
verify_model(Clone(), [([5, 3], "float32")], {}, Expected)
def test_lerp():
class Lerp(Module):
def forward(self, start, end, weight):
return torch.lerp(start, end, weight)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="float32"),
inp_1: R.Tensor((5, 3), dtype="float32"),
inp_2: R.Tensor((5, 3), dtype="float32"),
) -> R.Tensor((5, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((5, 3), dtype="float32") = R.add(
inp_0, R.multiply(inp_2, R.subtract(inp_1, inp_0))
)
gv: R.Tensor((5, 3), dtype="float32") = lv
R.output(gv)
return gv
verify_model(
Lerp(), [([5, 3], "float32"), ([5, 3], "float32"), ([5, 3], "float32")], {}, Expected
)
def test_std():
class Std(Module):
def forward(self, x):
return torch.std(x)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="float32"),
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((), dtype="float32") = R.std(inp_0, axis=None, keepdims=False)
gv: R.Tensor((), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Std(), [([5, 3], "float32")], {}, Expected)
def test_var():
class Var(Module):
def forward(self, x):
return torch.var(x)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="float32"),
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((), dtype="float32") = R.variance(inp_0, axis=None, keepdims=False)
gv: R.Tensor((), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Var(), [([5, 3], "float32")], {}, Expected)
@pytest.mark.parametrize(
"torch_dtype,relax_dtype",
[(torch.float32, "float32"), (torch.bool, "bool")],
)
def test_prod(torch_dtype, relax_dtype):
class Prod(Module):
def forward(self, x):
return torch.prod(x, dtype=torch_dtype)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype=relax_dtype),
) -> R.Tensor((), dtype=relax_dtype):
with R.dataflow():
lv: R.Tensor((), dtype=relax_dtype) = R.prod(inp_0, axis=None, keepdims=False)
gv: R.Tensor((), dtype=relax_dtype) = lv
R.output(gv)
return gv
verify_model(Prod(), [([5, 3], relax_dtype)], {}, Expected)
def test_cumprod():
class Cumprod(Module):
def forward(self, x):
return torch.cumprod(x, 0)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="float32"),
) -> R.Tensor((5, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((5, 3), dtype="float32") = R.cumprod(inp_0, axis=0, exclusive=False)
gv: R.Tensor((5, 3), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Cumprod(), [([5, 3], "float32")], {}, Expected)
def test_where():
class Where(Module):
def forward(self, condition, x, y):
return torch.where(condition, x, y)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="bool"),
inp_1: R.Tensor((5, 3), dtype="float32"),
inp_2: R.Tensor((5, 3), dtype="float32"),
) -> R.Tensor((5, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((5, 3), dtype="float32") = R.where(inp_0, inp_1, inp_2)
gv: R.Tensor((5, 3), dtype="float32") = lv
R.output(gv)
return gv
verify_model(
Where(), [([5, 3], "bool"), ([5, 3], "float32"), ([5, 3], "float32")], {}, Expected
)
def test_bucketize():
class Bucketize(Module):
def forward(self, input_tensor, boundaries):
return torch.bucketize(input_tensor, boundaries)
@tvm.script.ir_module
class Expected:
@R.function
def main(
input: R.Tensor((5, 3), dtype="float32"), boundaries: R.Tensor((10,), dtype="float32")
) -> R.Tensor((5, 3), dtype="int64"):
with R.dataflow():
lv: R.Tensor((5, 3), dtype="int64") = R.bucketize(
input, boundaries, out_int32=False, right=False
)
gv: R.Tensor((5, 3), dtype="int64") = lv
R.output(gv)
return gv
verify_model(Bucketize(), [([5, 3], "float32"), ([10], "float32")], {}, Expected)
def test_argsort():
class Argsort(Module):
def forward(self, x):
return torch.argsort(x, dim=1, descending=True)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="float32"),
) -> R.Tensor((5, 3), dtype="int32"):
with R.dataflow():
lv: R.Tensor((5, 3), dtype="int32") = R.argsort(inp_0, axis=1, descending=True)
gv: R.Tensor((5, 3), dtype="int32") = lv
R.output(gv)
return gv
verify_model(Argsort(), [([5, 3], "float32")], {}, Expected)
def test_sort():
class Sort(Module):
def forward(self, x):
return torch.sort(x, dim=1, descending=True)
@tvm.script.ir_module
class Expected:
@R.function
def main(inp_0: R.Tensor((5, 3), dtype="float32")) -> R.Tuple(
R.Tensor((5, 3), dtype="float32"), R.Tensor((5, 3), dtype="int32")
):
with R.dataflow():
lv: R.Tensor((5, 3), dtype="int32") = R.argsort(
inp_0, axis=1, descending=True, dtype="int32"
)
lv1: R.Tensor((5, 3), dtype="float32") = R.gather_elements(inp_0, lv, axis=1)
lv2: R.Tuple(R.Tensor((5, 3), dtype="float32"), R.Tensor((5, 3), dtype="int32")) = (
lv1,
lv,
)
gv: R.Tuple(R.Tensor((5, 3), dtype="float32"), R.Tensor((5, 3), dtype="int32")) = (
lv2
)
R.output(gv)
return gv
verify_model(Sort(), [([5, 3], "float32")], {}, Expected)
def test_topk():
class Topk(Module):
def forward(self, x):
return torch.topk(x, k=2, dim=1, largest=True, sorted=True)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="float32"),
) -> R.Tuple(R.Tensor((5, 2), dtype="float32"), R.Tensor((5, 2), dtype="int64")):
with R.dataflow():
lv: R.Tuple(R.Tensor((5, 2), dtype="float32"), R.Tensor((5, 2), dtype="int64")) = (
R.topk(inp_0, k=2, axis=1, ret_type="both", largest=True, dtype="int64")
)
gv: R.Tuple(R.Tensor((5, 2), dtype="float32"), R.Tensor((5, 2), dtype="int64")) = lv
R.output(gv)
return gv
verify_model(Topk(), [([5, 3], "float32")], {}, Expected)
def test_broadcast_to():
class BroadcastTo(Module):
def forward(self, x):
return torch.broadcast_to(x, (5, 3))
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 1), dtype="float32"),
) -> R.Tensor((5, 3), dtype="float32"):
with R.dataflow():
lv: R.Tensor((5, 3), dtype="float32") = R.broadcast_to(inp_0, (5, 3))
gv: R.Tensor((5, 3), dtype="float32") = lv
R.output(gv)
return gv
verify_model(BroadcastTo(), [([5, 1], "float32")], {}, Expected)
def test_narrow():
class Narrow(Module):
def forward(self, x):
return torch.narrow(x, 1, 0, 2)
@tvm.script.ir_module
class Expected:
@R.function
def main(
inp_0: R.Tensor((5, 3), dtype="float32"),
) -> R.Tensor((5, 2), dtype="float32"):
with R.dataflow():
lv: R.Tensor((5, 2), dtype="float32") = R.strided_slice(
inp_0, axes=[1], begin=[0], end=[2]
)
gv: R.Tensor((5, 2), dtype="float32") = lv
R.output(gv)
return gv
verify_model(Narrow(), [([5, 3], "float32")], {}, Expected)
def test_norm():
input_info = [([1, 3, 5, 3], "float32")]
class Norm(Module):
def __init__(self, p, dim=None, keepdim=False):
super().__init__()
self.p = p
self.dim = dim
self.keepdim = keepdim
def forward(self, x):
return torch.norm(x, p=self.p, dim=self.dim, keepdim=self.keepdim)
@tvm.script.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((), dtype="float32") = R.max(R.abs(inp_0), axis=None, keepdims=False)
gv: R.Tensor((), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class Expected2:
@R.function
def main(
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((), dtype="float32") = R.min(R.abs(inp_0), axis=None, keepdims=False)
gv: R.Tensor((), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class Expected3:
@R.function
def main(
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(2, "float32"))
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
lv3: R.Tensor((), dtype="float32") = R.power(lv2, R.const(0.5, "float32"))
gv: R.Tensor((), dtype="float32") = lv3
R.output(gv)
return gv
@tvm.script.ir_module
class Expected4:
@R.function
def main(
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(1.0, "float32"))
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
lv3: R.Tensor((), dtype="float32") = R.power(lv2, R.const(1.0, "float32"))
gv: R.Tensor((), dtype="float32") = lv3
R.output(gv)
return gv
@tvm.script.ir_module
class Expected5:
@R.function
def main(
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(-4, "float32"))
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
lv3: R.Tensor((), dtype="float32") = R.power(lv2, R.const(-0.25, "float32"))
gv: R.Tensor((), dtype="float32") = lv3
R.output(gv)
return gv
@tvm.script.ir_module
class Expected6:
@R.function
def main(
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.abs(inp_0)
lv1: R.Tensor((1, 3, 5, 3), dtype="float32") = R.power(lv, R.const(0.5, "float32"))
lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False)
lv3: R.Tensor((), dtype="float32") = R.power(lv2, R.const(2, "float32"))
gv: R.Tensor((), dtype="float32") = lv3
R.output(gv)
return gv
@tvm.script.ir_module
class Expected7:
@R.function
def main(
inp_0: R.Tensor((1, 3, 5, 3), dtype="float32"),
) -> R.Tensor((), dtype="float32"):
with R.dataflow():
lv: R.Tensor((1, 3, 5, 3), dtype="float32") = R.multiply(inp_0, inp_0)
lv1: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
lv2: R.Tensor((), dtype="float32") = R.sqrt(lv1)
gv: R.Tensor((), dtype="float32") = lv2
R.output(gv)
return gv
norms = [
((float("inf"), None, False), Expected1),
((float("-inf"), None, False), Expected2),
((float(2), None, False), Expected3),
((1.0, None, False), Expected4),
((float(-4), None, True), Expected5),
((0.5, None, True), Expected6),
(("fro", None, False), Expected7),
]
for (p, dim, keepdim), expected in norms:
verify_model(Norm(p, dim=dim, keepdim=keepdim), input_info, {}, expected)
@pytest.mark.parametrize(
"torch_dtype, relax_dtype",
[
# Float types
(torch.float16, "float16"),
(torch.float32, "float32"),
(torch.float64, "float64"),
(torch.bfloat16, "bfloat16"),
# Signed integer types
(torch.int8, "int8"),
(torch.int16, "int16"),
(torch.int32, "int32"),
(torch.int64, "int64"),
# Unsigned integer types
(torch.uint8, "uint8"),
(torch.uint16, "uint16"),
(torch.uint32, "uint32"),
(torch.uint64, "uint64"),
# Boolean
(torch.bool, "bool"),
],
)
def test_dtypes(torch_dtype, relax_dtype):
class Model(Module):
def forward(self, lhs: torch.Tensor, rhs: torch.Tensor):
return torch.ops.aten.add(lhs, rhs)
@tvm.script.ir_module
class Expected:
@R.function
def main(
lhs: R.Tensor((10, 10), dtype=relax_dtype),
rhs: R.Tensor((10, 10), dtype=relax_dtype),
) -> R.Tensor((10, 10), dtype=relax_dtype):
with R.dataflow():
lv: R.Tensor((10, 10), dtype=relax_dtype) = relax.op.add(lhs, rhs)
gv: R.Tensor((10, 10), dtype=relax_dtype) = lv
R.output(gv)
return gv
verify_model(Model(), [([10, 10], torch_dtype), ([10, 10], torch_dtype)], {}, Expected)
def test_eye():
import numpy as np
class Eye(Module):
def forward(self, input):
return torch.eye(3)
graph_model = fx.symbolic_trace(Eye())
mod = from_fx(graph_model, [([3, 3], "float32")])
assert len(mod["main"].body.blocks) == 1
assert len(mod["main"].body.blocks[0].bindings) == 1
assert isinstance(mod["main"].body.blocks[0].bindings[0].value, relax.Constant)
tvm.testing.assert_allclose(
mod["main"].body.blocks[0].bindings[0].value.data.numpy(),
np.eye(3, dtype="float32"),
)
def test_linspace():
import numpy as np
class Linspace(Module):
def forward(self, input):
return torch.linspace(0, 1, steps=9)
graph_model = fx.symbolic_trace(Linspace())
mod = from_fx(graph_model, [([9, 9], "float32")])
assert len(mod["main"].body.blocks) == 1
assert len(mod["main"].body.blocks[0].bindings) == 1
assert isinstance(mod["main"].body.blocks[0].bindings[0].value, relax.Constant)
tvm.testing.assert_allclose(
mod["main"].body.blocks[0].bindings[0].value.data.numpy(),
np.linspace(0, 1, num=9, dtype="float32"),
)
def test_round():
input_info = [([3, 4], "float32")]
class Round(Module):
def __init__(self, decimals=0):
super().__init__()
self.decimals = decimals
def forward(self, x):
if self.decimals == 0:
return torch.round(x)
else:
return torch.round(x, decimals=self.decimals)
@tvm.script.ir_module
class Expected1:
@R.function
def main(
inp_0: R.Tensor((3, 4), dtype="float32"),
) -> R.Tensor((3, 4), dtype="float32"):
with R.dataflow():
lv: R.Tensor((3, 4), dtype="float32") = R.round(inp_0)
gv: R.Tensor((3, 4), dtype="float32") = lv
R.output(gv)
return gv
@tvm.script.ir_module
class Expected2:
@R.function
def main(
inp_0: R.Tensor((3, 4), dtype="float32"),
) -> R.Tensor((3, 4), dtype="float32"):
with R.dataflow():
lv: R.Tensor((3, 4), dtype="float32") = R.multiply(inp_0, R.const(100.0, "float32"))
lv1: R.Tensor((3, 4), dtype="float32") = R.round(lv)
lv2: R.Tensor((3, 4), dtype="float32") = R.divide(lv1, R.const(100.0, "float32"))
gv: R.Tensor((3, 4), dtype="float32") = lv2
R.output(gv)
return gv
rounds = [
(0, Expected1),
(2, Expected2),
]
for decimals, expected in rounds:
verify_model(Round(decimals), input_info, {}, expected)
# Test numerical accuracy with decimals
test_data = torch.tensor(
[
[1.2345, 2.3456, 3.4567, 4.5678],
[5.6789, 6.7890, 7.8901, 8.9012],
[9.1234, 10.2345, 11.3456, 12.4567],
]
)
for decimals in [0, 1, 2, 3]:
torch_model = Round(decimals)
graph_model = fx.symbolic_trace(torch_model)
with torch.no_grad():
mod = from_fx(graph_model, input_info)
target = tvm.target.Target("llvm")
ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
torch_result = torch_model(test_data).numpy()
tvm_input = tvm.runtime.tensor(test_data.numpy())
tvm_result = vm["main"](tvm_input).numpy()
# Use relaxed tolerance due to floating-point precision in decimal operations
tvm.testing.assert_allclose(tvm_result, torch_result, rtol=1e-3, atol=1e-3)
if __name__ == "__main__":
tvm.testing.main()
Reference in New Issue View Git Blame Copy Permalink