This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 9f1e1980c1 [Tests][Frontend] Remove redundant PyTorch frontend tests
(#20021)
9f1e1980c1 is described below
commit 9f1e1980c1442652df69d64f506c5e5e70efdebe
Author: Shushi Hong <[email protected]>
AuthorDate: Thu Jul 16 23:06:12 2026 -0400
[Tests][Frontend] Remove redundant PyTorch frontend tests (#20021)
This PR:
- Removes duplicate module/functional, alias, positional-argument, and
no-op cases from the PyTorch ExportedProgram tests.
- Consolidates the four GRU configurations into a table-driven loop
without removing any configurations.
- Removes duplicated FX cases already covered through the same shared
converters.
- Restores FX scalar tensor constant coverage and TFLite
constant-parameter Gather and static broadcast/MUL coverage.
---
.../relax/test_frontend_from_exported_program.py | 715 ++-------------------
tests/python/relax/test_frontend_from_fx.py | 468 +-------------
tests/python/relax/test_frontend_tflite.py | 69 ++
3 files changed, 147 insertions(+), 1105 deletions(-)
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index 1f78e941e4..3799b8ed95 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: E501, F401, F841
+# ruff: noqa: F401, F841
import operator
import numpy as np
@@ -222,15 +222,7 @@ def test_extended_unary_ops():
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.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):
+ class Celu(Module):
def forward(self, input):
return torch.nn.functional.celu(input)
@@ -254,8 +246,7 @@ def test_extended_unary_ops():
R.output(gv)
return gv
- verify_model(Celu1(), example_args, {}, expected_celu)
- verify_model(Celu2(), example_args, {}, expected_celu)
+ verify_model(Celu(), example_args, {}, expected_celu)
# clamp
class Clamp(Module):
@@ -330,15 +321,7 @@ def test_extended_unary_ops():
# 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):
+ class Dropout(Module):
def forward(self, input):
return torch.dropout(input, 0.5, train=True)
@@ -377,20 +360,11 @@ def test_extended_unary_ops():
R.output(gv)
return gv
- verify_model(Dropout1(), example_args, {}, expected_dropout_for_1_2)
- verify_model(Dropout2(), example_args, {}, expected_dropout_for_1_2)
+ verify_model(Dropout(), example_args, {}, expected_dropout_for_1_2)
verify_model(Dropout3(), example_args, {}, expected_dropout_for_3)
# 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)
@@ -416,18 +390,9 @@ def test_extended_unary_ops():
return gv
verify_model(Elu(), example_args, {}, expected_elu)
- verify_model(Elu2(), example_args, {}, expected_elu)
# 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)
@@ -455,18 +420,9 @@ def test_extended_unary_ops():
return gv
verify_model(Hardsigmoid(), example_args, {}, expected_hardsigmoid)
- verify_model(Hardsigmoid2(), example_args, {}, expected_hardsigmoid)
# hardwish
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)
@@ -523,7 +479,6 @@ def test_extended_unary_ops():
return gv
verify_model(Hardswish(), example_args, {}, expected_hardswish_for_1_2)
- verify_model(Hardswish2(), example_args, {}, expected_hardswish_for_1_2)
verify_model(Hardswish3(), example_args, {}, expected_hardswish_for_3)
# isfinite
@@ -846,14 +801,6 @@ def test_extended_unary_ops():
def test_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)
@@ -877,26 +824,12 @@ def test_hardtanh():
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
verify_model(Hardtanh(), example_args, {}, expected_for_1_2)
- verify_model(Hardtanh2(), example_args, {}, expected_for_1_2)
# In-place hardtanh_ yields the same program; mutation outputs are dropped.
verify_model(Hardtanh3(), example_args, {}, expected_for_1_2)
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):
+ class Softplus(Module):
def forward(self, input):
return torch.nn.functional.softplus(input, 1.0, 20.0)
@@ -925,25 +858,11 @@ def test_softplus():
return gv
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
- verify_model(Softplus0(), example_args, {}, expected)
- verify_model(Softplus1(), example_args, {}, expected)
+ verify_model(Softplus(), example_args, {}, expected)
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):
+ class LeakyReLU(Module):
def forward(self, input):
return torch.nn.functional.leaky_relu(input, 0.02)
@@ -965,8 +884,7 @@ def test_leakyrelu():
return gv
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
- verify_model(LeakyReLU0(), example_args, {}, expected_for_1_2)
- verify_model(LeakyReLU1(), example_args, {}, expected_for_1_2)
+ verify_model(LeakyReLU(), example_args, {}, expected_for_1_2)
# In-place leaky_relu_ yields the same program; mutation outputs are
dropped.
verify_model(LeakyReLU2(), example_args, {}, expected_for_1_2)
@@ -1173,14 +1091,6 @@ def test_pow_integer():
def test_logsoftmax():
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)
@@ -1199,19 +1109,10 @@ def test_logsoftmax():
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
verify_model(LogSoftmax(), example_args, {}, expected1)
- verify_model(LogSoftmax2(), example_args, {}, expected1)
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):
+ class Prelu(torch.nn.Module):
def __init__(self):
super().__init__()
self.alpha = torch.nn.Parameter(torch.tensor([0.25]))
@@ -1237,20 +1138,11 @@ def test_prelu():
return gv
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
- verify_model(Prelu1(), example_args, {}, expected)
- verify_model(Prelu2(), example_args, {}, expected)
+ verify_model(Prelu(), example_args, {}, expected)
def test_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)
@@ -1269,19 +1161,10 @@ def test_softmax():
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
verify_model(Softmax(), example_args, {}, expected1)
- verify_model(Softmax2(), example_args, {}, expected1)
def test_softsign():
class Softsign(Module):
- def __init__(self):
- super().__init__()
- self.ss = torch.nn.Softsign()
-
- def forward(self, input):
- return self.ss(input)
-
- class Softsign2(Module):
def forward(self, input):
return torch.nn.functional.softsign(input)
@@ -1301,19 +1184,10 @@ def test_softsign():
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
verify_model(Softsign(), example_args, {}, expected_softsign)
- verify_model(Softsign2(), example_args, {}, expected_softsign)
def test_softshrink():
class Softshrink(Module):
- def __init__(self):
- super().__init__()
- self.softshrink = torch.nn.Softshrink(lambd=0.5)
-
- def forward(self, input):
- return self.softshrink(input)
-
- class Softshrink2(Module):
def forward(self, input):
return torch.nn.functional.softshrink(input, lambd=0.5)
@@ -1341,7 +1215,6 @@ def test_softshrink():
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
verify_model(Softshrink(), example_args, {}, expected_softshrink)
- verify_model(Softshrink2(), example_args, {}, expected_softshrink)
def test_tril_triu():
@@ -1411,7 +1284,6 @@ def test_tril_triu():
operator_binary_1 = [
(operator.add, R.add),
(torch.ops.aten.add_, R.add),
- (torch.ops.aten.bitwise_or, R.bitwise_or),
(torch.ops.aten.bitwise_or_, R.bitwise_or),
(operator.sub, R.subtract),
(operator.mul, R.multiply),
@@ -1996,15 +1868,7 @@ def test_batchnorm2d():
def test_adaptive_avgpool1d():
- class AdaptiveAvgPool1d0(torch.nn.Module):
- def __init__(self):
- super().__init__()
- self.pool = torch.nn.AdaptiveAvgPool1d(output_size=5)
-
- def forward(self, input):
- return self.pool(input)
-
- class AdaptiveAvgPool1d1(torch.nn.Module):
+ class AdaptiveAvgPool1d(torch.nn.Module):
def forward(self, input):
return torch.nn.functional.adaptive_avg_pool1d(input,
output_size=5)
@@ -2025,20 +1889,11 @@ def test_adaptive_avgpool1d():
return gv
example_args = (torch.randn(1, 3, 10, dtype=torch.float32),)
- verify_model(AdaptiveAvgPool1d0(), example_args, {}, expected1)
- verify_model(AdaptiveAvgPool1d1(), example_args, {}, expected1)
+ verify_model(AdaptiveAvgPool1d(), example_args, {}, expected1)
def test_adaptive_avgpool2d():
- 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):
+ class AdaptiveAvgPool2d(Module):
def forward(self, input):
return torch.nn.functional.adaptive_avg_pool2d(input, [10, 10])
@@ -2058,20 +1913,11 @@ def test_adaptive_avgpool2d():
return gv
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
- verify_model(AdaptiveAvgPool2d0(), example_args, {}, expected1)
- verify_model(AdaptiveAvgPool2d1(), example_args, {}, expected1)
+ verify_model(AdaptiveAvgPool2d(), example_args, {}, expected1)
def test_adaptive_avgpool3d():
- class AdaptiveAvgPool3d0(torch.nn.Module):
- def __init__(self):
- super().__init__()
- self.pool = torch.nn.AdaptiveAvgPool3d([4, 4, 4])
-
- def forward(self, input):
- return self.pool(input)
-
- class AdaptiveAvgPool3d1(torch.nn.Module):
+ class AdaptiveAvgPool3d(torch.nn.Module):
def forward(self, input):
return torch.nn.functional.adaptive_avg_pool3d(input, [4, 4, 4])
@@ -2090,8 +1936,7 @@ def test_adaptive_avgpool3d():
return gv
example_args = (torch.randn(1, 3, 8, 8, 8, dtype=torch.float32),)
- verify_model(AdaptiveAvgPool3d0(), example_args, {}, expected1)
- verify_model(AdaptiveAvgPool3d1(), example_args, {}, expected1)
+ verify_model(AdaptiveAvgPool3d(), example_args, {}, expected1)
def test_addmm():
@@ -2244,14 +2089,6 @@ def test_avg_pool1d():
return gv
class AvgPool1d2(Module):
- def __init__(self):
- super().__init__()
- self.pool = torch.nn.AvgPool1d(kernel_size=3, stride=2, padding=1,
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=3, stride=2, padding=1, ceil_mode=True
@@ -2312,7 +2149,6 @@ def test_avg_pool1d():
example_args = (torch.randn(1, 3, 10, dtype=torch.float32),)
verify_model(AvgPool1d1(), example_args, {}, expected1)
verify_model(AvgPool1d2(), example_args, {}, expected2)
- verify_model(AvgPool1d3(), example_args, {}, expected2)
verify_model(AvgPool1d4(), example_args, {}, expected3)
@@ -2348,14 +2184,6 @@ def test_avg_pool2d():
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
@@ -2434,7 +2262,6 @@ def test_avg_pool2d():
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
verify_model(AvgPool2d1(), example_args, {}, expected1)
verify_model(AvgPool2d2(), example_args, {}, expected2)
- verify_model(AvgPool2d3(), example_args, {}, expected2)
verify_model(AvgPool2d4(), example_args, {}, expected4)
verify_model(AvgPool2d5(), example_args, {}, expected5)
@@ -2471,14 +2298,6 @@ def test_avg_pool3d():
return gv
class AvgPool3d2(Module):
- def __init__(self):
- super().__init__()
- self.pool = torch.nn.AvgPool3d(kernel_size=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, stride=2, padding=1, ceil_mode=True
@@ -2531,7 +2350,6 @@ def test_avg_pool3d():
example_args = (torch.randn(1, 3, 8, 8, 8, dtype=torch.float32),)
verify_model(AvgPool3d1(), example_args, {}, expected1)
verify_model(AvgPool3d2(), example_args, {}, expected2)
- verify_model(AvgPool3d3(), example_args, {}, expected2)
verify_model(AvgPool3d4(), example_args, {}, expected3)
@@ -3429,15 +3247,7 @@ def test_pad():
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):
+ class PixelShuffle(torch.nn.Module):
def __init__(self, upscale_factor=2):
super().__init__()
self.upscale_factor = upscale_factor
@@ -3466,8 +3276,7 @@ def test_pixel_shuffle():
return gv
example_args = (torch.randn(1, 8, 10, 15, dtype=torch.float32),)
- verify_model(PixelShuffle1(upscale_factor=2), example_args, {}, expected)
- verify_model(PixelShuffle2(upscale_factor=2), example_args, {}, expected)
+ verify_model(PixelShuffle(upscale_factor=2), example_args, {}, expected)
def test_einsum():
@@ -3574,12 +3383,6 @@ def test_embedding():
def test_groupnorm():
- import torch
- from torch.nn import Module
-
- torch.set_grad_enabled(False)
- torch.random.manual_seed(0)
-
class GroupNorm(Module):
def __init__(self):
super().__init__()
@@ -3623,9 +3426,6 @@ def test_groupnorm():
def test_instancenorm2d():
- torch.set_grad_enabled(False)
- torch.random.manual_seed(0)
-
class InstanceNorm2d(Module):
def __init__(self):
super().__init__()
@@ -3792,14 +3592,6 @@ def test_linear():
def test_maxpool1d():
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__()
@@ -3842,34 +3634,6 @@ def test_maxpool1d():
R.output(gv)
return gv
- @tvm.script.ir_module
- class expected2:
- @R.function
- def main(input_1: R.Tensor((1, 3, 8), dtype="float32")) -> R.Tuple(
- R.Tensor((1, 3, 4), dtype="float32")
- ):
- with R.dataflow():
- lv: R.Tensor((1, 3, 1, 8), dtype="float32") =
R.expand_dims(input_1, axis=[-2])
- lv1: R.Tensor((1, 3, 1, 4), dtype="float32") = R.nn.max_pool2d(
- lv,
- pool_size=[1, 2],
- strides=[1, 2],
- dilation=[1, 1],
- padding=[0, 0, 0, 0],
- layout="NCHW",
- out_layout="NCHW",
- )
- lv2: R.Tensor((1, 3, 1, 4), dtype="float32") =
R.zeros_like(lv1)
- lv3: R.Tuple(
- R.Tensor((1, 3, 1, 4), dtype="float32"),
- R.Tensor((1, 3, 1, 4), dtype="float32"),
- ) = (lv1, lv2)
- lv4: R.Tensor((1, 3, 1, 4), dtype="float32") = lv3[0]
- lv5: R.Tensor((1, 3, 4), dtype="float32") = R.squeeze(lv4,
axis=[-2])
- gv: R.Tuple(R.Tensor((1, 3, 4), dtype="float32")) = (lv5,)
- R.output(gv)
- return gv
-
@tvm.script.ir_module
class expected3:
@R.function
@@ -3900,12 +3664,10 @@ def test_maxpool1d():
# Example inputs
example_args1 = (torch.randn(1, 3, 8, dtype=torch.float32),)
- example_args2 = (torch.randn(1, 3, 8, dtype=torch.float32),)
example_args3 = (torch.randn(1, 3, 10, dtype=torch.float32),)
# Verify the models
verify_model(MaxPool1d(), example_args1, {}, expected1)
- verify_model(MaxPool1d_functional(), example_args2, {}, expected2)
verify_model(MaxPool1d2(), example_args3, {}, expected3)
@@ -5195,7 +4957,7 @@ def test_argmax_argmin():
verify_model(Argmin2(), example_args, {}, expected_argmin2)
-def test_cat_concat():
+def test_cat():
class Cat0(Module):
def forward(self, x, y):
return torch.cat((x, y))
@@ -5204,14 +4966,6 @@ def test_cat_concat():
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
@@ -5241,8 +4995,6 @@ def test_cat_concat():
example_args = (torch.randn(2, 3, dtype=torch.float32), torch.randn(2, 3,
dtype=torch.float32))
verify_model(Cat0(), example_args, {}, Expected1)
verify_model(Cat1(), example_args, {}, Expected2)
- verify_model(Cat2(), example_args, {}, Expected2)
- verify_model(Cat3(), example_args, {}, Expected1)
def test_cumsum():
@@ -5376,10 +5128,6 @@ def test_permute():
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
@@ -5395,7 +5143,6 @@ def test_permute():
example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
verify_model(Permute1(), example_args, {}, expected1)
- verify_model(Permute2(), example_args, {}, expected1)
def test_repeat():
@@ -5437,9 +5184,6 @@ def test_repeat():
example_args = (torch.randn(1, 3, dtype=torch.float32),)
verify_model(Tile2(), example_args, {}, expected2)
- example_args = (torch.randn(1, 3, dtype=torch.float32),)
- verify_model(Tile2(), example_args, {}, expected2)
-
def test_reshape():
class Reshape(Module):
@@ -5776,7 +5520,6 @@ def test_split():
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"),
@@ -5794,111 +5537,9 @@ def test_split():
R.output(gv)
return gv
- class Unbind1(Module):
- def forward(self, data):
- return torch.unbind(data)
-
- @tvm.script.ir_module
- class expected1:
- @R.function
- def main(data: 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.Tensor((1, 3, 10, 10), dtype="float32") =
R.strided_slice(
- data,
- (R.prim_value(0),),
- (R.prim_value(0),),
- (R.prim_value(1),),
- (R.prim_value(1),),
- assume_inbound=False,
- )
- lv1: R.Tensor((1, 3, 10, 10), dtype="float32") =
R.strided_slice(
- data,
- (R.prim_value(0),),
- (R.prim_value(1),),
- (R.prim_value(2),),
- (R.prim_value(1),),
- assume_inbound=False,
- )
- lv2: R.Tensor((1, 3, 10, 10), dtype="float32") =
R.strided_slice(
- data,
- (R.prim_value(0),),
- (R.prim_value(2),),
- (R.prim_value(3),),
- (R.prim_value(1),),
- assume_inbound=False,
- )
- lv3: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv,
axis=[0])
- lv4: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv1,
axis=[0])
- lv5: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv2,
axis=[0])
- gv: R.Tuple(
- R.Tensor((3, 10, 10), dtype="float32"),
- R.Tensor((3, 10, 10), dtype="float32"),
- R.Tensor((3, 10, 10), dtype="float32"),
- ) = (lv3, lv4, lv5)
- R.output(gv)
- return gv
-
- class Unbind2(Module):
- def forward(self, data):
- return torch.unbind(data, dim=1)
-
- @tvm.script.ir_module
- class expected2:
- @R.function
- def main(data: 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.Tensor((3, 1, 10, 10), dtype="float32") =
R.strided_slice(
- data,
- (R.prim_value(1),),
- (R.prim_value(0),),
- (R.prim_value(1),),
- (R.prim_value(1),),
- assume_inbound=False,
- )
- lv1: R.Tensor((3, 1, 10, 10), dtype="float32") =
R.strided_slice(
- data,
- (R.prim_value(1),),
- (R.prim_value(1),),
- (R.prim_value(2),),
- (R.prim_value(1),),
- assume_inbound=False,
- )
- lv2: R.Tensor((3, 1, 10, 10), dtype="float32") =
R.strided_slice(
- data,
- (R.prim_value(1),),
- (R.prim_value(2),),
- (R.prim_value(3),),
- (R.prim_value(1),),
- assume_inbound=False,
- )
- lv3: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv,
axis=[1])
- lv4: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv1,
axis=[1])
- lv5: R.Tensor((3, 10, 10), dtype="float32") = R.squeeze(lv2,
axis=[1])
- gv: R.Tuple(
- R.Tensor((3, 10, 10), dtype="float32"),
- R.Tensor((3, 10, 10), dtype="float32"),
- R.Tensor((3, 10, 10), dtype="float32"),
- ) = (lv3, lv4, lv5)
- R.output(gv)
- return gv
-
example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
verify_model(Chunk(), example_args, {}, Expected)
- example_args = (torch.randn(3, 3, 10, 10, dtype=torch.float32),)
- verify_model(Unbind1(), example_args, {}, expected1)
- verify_model(Unbind2(), example_args, {}, expected2)
-
def test_squeeze():
class Squeeze1(Module):
@@ -5965,10 +5606,6 @@ def test_stack():
def forward(self, x, y):
return torch.stack((x, y), dim=1)
- class Stack2(Module):
- def forward(self, x, y):
- return torch.stack((x, y), 1) # positional dim
-
class Stack3(Module):
def forward(self, x, y):
return torch.stack((x, y), dim=-1) # negative dim
@@ -6020,7 +5657,6 @@ def test_stack():
verify_model(Stack0(), example_args, {}, Expected0)
verify_model(Stack1(), example_args, {}, Expected1)
- verify_model(Stack2(), example_args, {}, Expected1)
verify_model(Stack3(), example_args, {}, Expected3)
@@ -6033,10 +5669,6 @@ def test_tile():
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
@@ -6066,7 +5698,6 @@ def test_tile():
example_args = (torch.randn(1, 3, dtype=torch.float32),)
verify_model(Tile1(), example_args, {}, expected1)
verify_model(Tile2(), example_args, {}, expected2)
- verify_model(Tile3(), example_args, {}, expected2)
def test_transpose():
@@ -6240,26 +5871,6 @@ def test_hamming_window():
verify_model(HammingWindow(), example_args, {}, Expected)
-def test_contiguous():
- class Contiguous(Module):
- def forward(self, input):
- return input.contiguous()
-
- @tvm.script.ir_module
- class Expected:
- @R.function
- def main(
- input: R.Tensor((10, 10), dtype="float32"),
- ) -> R.Tuple(R.Tensor((10, 10), dtype="float32")):
- with R.dataflow():
- gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (input,)
- R.output(gv)
- return gv
-
- example_args = (torch.randn(10, 10, dtype=torch.float32),)
- verify_model(Contiguous(), example_args, {}, Expected)
-
-
def test_clone():
class Clone(Module):
def forward(self, input):
@@ -6538,95 +6149,24 @@ def test_copy():
def test_to_copy():
- # float
- class ToFloat(Module):
- def forward(self, x):
- return x.float()
-
- @tvm.script.ir_module
- class expected_float:
- @R.function
- def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
- R.Tensor((1, 2, 3, 4), dtype="float32")
- ):
- # block 0
- with R.dataflow():
- gv: R.Tuple(R.Tensor((1, 2, 3, 4), dtype="float32")) = (x,)
- R.output(gv)
- return gv
-
- # half
class ToHalf(Module):
def forward(self, x):
return x.half()
@tvm.script.ir_module
- class expected_half:
+ class Expected:
@R.function
def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
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.Tuple(R.Tensor((1, 2, 3, 4), dtype="float16")) = (lv,)
R.output(gv)
return gv
- # type
- class Type(Module):
- def forward(self, x):
- return x.type(torch.float32)
-
- @tvm.script.ir_module
- class expected_type:
- @R.function
- def main(x: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
- R.Tensor((1, 2, 3, 4), dtype="float32")
- ):
- # block 0
- with R.dataflow():
- gv: R.Tuple(R.Tensor((1, 2, 3, 4), dtype="float32")) = (x,)
- R.output(gv)
- return gv
-
- class To1(Module):
- def forward(self, input):
- return input.to(torch.float16)
-
- @I.ir_module
- class expected_to1:
- @R.function
- def main(input: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
- R.Tensor((1, 2, 3, 4), dtype="float16")
- ):
- with R.dataflow():
- lv: R.Tensor((1, 2, 3, 4), dtype="float16") = R.astype(input,
dtype="float16")
- gv: R.Tuple(R.Tensor((1, 2, 3, 4), dtype="float16")) = (lv,)
- R.output(gv)
- return gv
-
- class To2(Module):
- def forward(self, input):
- return input.to("cpu")
-
- @I.ir_module
- class expected_to2:
- @R.function
- def main(input: R.Tensor((1, 2, 3, 4), dtype="float32")) -> R.Tuple(
- R.Tensor((1, 2, 3, 4), dtype="float32")
- ):
- with R.dataflow():
- gv: R.Tuple(R.Tensor((1, 2, 3, 4), dtype="float32")) = (input,)
- R.output(gv)
- return gv
-
example_args = (torch.randn(1, 2, 3, 4, dtype=torch.float32),)
- verify_model(ToFloat(), example_args, {}, expected_float)
- verify_model(ToHalf(), example_args, {}, expected_half)
- verify_model(Type(), example_args, {}, expected_type)
- verify_model(To1(), example_args, {}, expected_to1)
- verify_model(To2(), example_args, {}, expected_to2)
+ verify_model(ToHalf(), example_args, {}, Expected)
def test_keep_params():
@@ -8431,160 +7971,36 @@ def test_tensor_none_tuple():
def test_gru():
- class BasicGRU(nn.Module):
- def __init__(self):
- super().__init__()
- self.gru = nn.GRU(
- input_size=4,
- hidden_size=8,
- num_layers=1,
- batch_first=True,
- bidirectional=False,
- )
-
- def forward(self, x):
- y, _ = self.gru(x)
- return y
-
- torch.manual_seed(42)
- x = torch.randn(2, 3, 4, dtype=torch.float32)
- model = BasicGRU()
- with torch.no_grad():
- pytorch_output = model(x)
- exported_program = export(model, args=(x,))
- mod = from_exported_program(exported_program)
- target = tvm.target.Target("llvm")
- ex = relax.build(mod, target)
- vm = relax.VirtualMachine(ex, tvm.cpu())
- x_tvm = tvm.runtime.tensor(x.numpy())
- tvm_output = vm["main"](x_tvm)
- if hasattr(tvm_output, "numpy"):
- tvm_output_np = tvm_output.numpy()
- else:
- tvm_output_np = tvm_output[0].numpy()
- assert pytorch_output.shape == tvm_output_np.shape, (
- f"Shape mismatch: PyTorch {pytorch_output.shape} vs TVM
{tvm_output_np.shape}"
- )
- tvm.testing.assert_allclose(pytorch_output.numpy(), tvm_output_np,
rtol=1e-4, atol=1e-5)
-
- class SeqFirstGRU(nn.Module):
- def __init__(self):
- super().__init__()
- self.gru = nn.GRU(
- input_size=3,
- hidden_size=6,
- num_layers=1,
- batch_first=False,
- bidirectional=False,
- )
-
- def forward(self, x):
- y, _ = self.gru(x)
- return y
-
- torch.manual_seed(43)
- x2 = torch.randn(4, 2, 3, dtype=torch.float32)
- model2 = SeqFirstGRU()
- with torch.no_grad():
- pytorch_output2 = model2(x2)
- exported_program2 = export(model2, args=(x2,))
- mod2 = from_exported_program(exported_program2)
- ex2 = relax.build(mod2, target)
- vm2 = relax.VirtualMachine(ex2, tvm.cpu())
- x2_tvm = tvm.runtime.tensor(x2.numpy())
- tvm_output2 = vm2["main"](x2_tvm)
- if hasattr(tvm_output2, "numpy"):
- tvm_output2_np = tvm_output2.numpy()
- else:
- tvm_output2_np = tvm_output2[0].numpy()
- assert pytorch_output2.shape == tvm_output2_np.shape
- tvm.testing.assert_allclose(pytorch_output2.numpy(), tvm_output2_np,
rtol=1e-4, atol=1e-5)
-
- # Test bidirectional GRU with batch_first=True
- class BidirectionalGRU(nn.Module):
- def __init__(self):
- super().__init__()
- self.gru = nn.GRU(
- input_size=4,
- hidden_size=5,
- num_layers=1,
- batch_first=True,
- bidirectional=True,
- )
-
- def forward(self, x):
- y, _ = self.gru(x)
- return y
-
- torch.manual_seed(44)
- x3 = torch.randn(2, 3, 4, dtype=torch.float32)
- model3 = BidirectionalGRU()
- with torch.no_grad():
- pytorch_output3 = model3(x3)
-
- # Verify output shape is correct (hidden_size * 2 due to bidirectional)
- assert pytorch_output3.shape == (
- 2,
- 3,
- 10,
- ), f"Expected shape (2, 3, 10), got {pytorch_output3.shape}"
-
- exported_program3 = export(model3, args=(x3,))
- mod3 = from_exported_program(exported_program3)
- ex3 = relax.build(mod3, target)
- vm3 = relax.VirtualMachine(ex3, tvm.cpu())
- x3_tvm = tvm.runtime.tensor(x3.numpy())
- tvm_output3 = vm3["main"](x3_tvm)
- if hasattr(tvm_output3, "numpy"):
- tvm_output3_np = tvm_output3.numpy()
- else:
- tvm_output3_np = tvm_output3[0].numpy()
- assert pytorch_output3.shape == tvm_output3_np.shape, (
- f"Shape mismatch: PyTorch {pytorch_output3.shape} vs TVM
{tvm_output3_np.shape}"
- )
- tvm.testing.assert_allclose(pytorch_output3.numpy(), tvm_output3_np,
rtol=1e-4, atol=1e-5)
-
- # Test bidirectional GRU with batch_first=False
- class SeqFirstBidirectionalGRU(nn.Module):
- def __init__(self):
+ class GRU(nn.Module):
+ def __init__(self, input_size, hidden_size, batch_first,
bidirectional):
super().__init__()
self.gru = nn.GRU(
- input_size=3,
- hidden_size=4,
+ input_size=input_size,
+ hidden_size=hidden_size,
num_layers=1,
- batch_first=False,
- bidirectional=True,
+ batch_first=batch_first,
+ bidirectional=bidirectional,
)
def forward(self, x):
y, _ = self.gru(x)
return y
- torch.manual_seed(45)
- x4 = torch.randn(4, 2, 3, dtype=torch.float32) # (seq_len, batch,
input_size)
- model4 = SeqFirstBidirectionalGRU()
- with torch.no_grad():
- pytorch_output4 = model4(x4)
-
- # Verify output shape (seq_len, batch, hidden_size * 2)
- assert pytorch_output4.shape == (
- 4,
- 2,
- 8,
- ), f"Expected shape (4, 2, 8), got {pytorch_output4.shape}"
-
- exported_program4 = export(model4, args=(x4,))
- mod4 = from_exported_program(exported_program4)
- ex4 = relax.build(mod4, target)
- vm4 = relax.VirtualMachine(ex4, tvm.cpu())
- x4_tvm = tvm.runtime.tensor(x4.numpy())
- tvm_output4 = vm4["main"](x4_tvm)
- if hasattr(tvm_output4, "numpy"):
- tvm_output4_np = tvm_output4.numpy()
- else:
- tvm_output4_np = tvm_output4[0].numpy()
- assert pytorch_output4.shape == tvm_output4_np.shape
- tvm.testing.assert_allclose(pytorch_output4.numpy(), tvm_output4_np,
rtol=1e-4, atol=1e-5)
+ cases = [
+ (42, (2, 3, 4), 4, 8, True, False),
+ (43, (4, 2, 3), 3, 6, False, False),
+ (44, (2, 3, 4), 4, 5, True, True),
+ (45, (4, 2, 3), 3, 4, False, True),
+ ]
+ for seed, shape, input_size, hidden_size, batch_first, bidirectional in
cases:
+ torch.manual_seed(seed)
+ x = torch.randn(*shape, dtype=torch.float32)
+ verify_model_numerically(
+ GRU(input_size, hidden_size, batch_first, bidirectional),
+ (x,),
+ rtol=1e-4,
+ atol=1e-5,
+ )
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -8814,45 +8230,12 @@ def test_dynamic_shape_with_unbounded_constraints():
def test_sym_size_int():
class SymSizeInt(Module):
- def __init__(self, dim):
- super().__init__()
- self.dim = dim
-
def forward(self, x):
- # TODO(@mshr-h): `torch.ops.aten.sym_size.int(x, self.dim)` would
be ideal, but currently
- # the ep frontend is not able to handle it.
- return torch.add(x[0], torch.ops.aten.sym_size.int(x, self.dim))
-
- @I.ir_module
- class Expected1:
- @R.function
- def main(x: R.Tensor((1, 3, 4), dtype="float32")) -> R.Tuple(
- R.Tensor((3, 4), dtype="float32")
- ):
- with R.dataflow():
- lv: R.Tensor((3, 4), dtype="float32") = R.take(
- x, R.const(0, "int64"), axis=0, mode="fast"
- )
- lv1: R.Tensor((3, 4), dtype="float32") = R.add(lv,
R.const(3.0, "float32"))
- gv: R.Tuple(R.Tensor((3, 4), dtype="float32")) = (lv1,)
- R.output(gv)
- return gv
-
- example_args_1 = (torch.randn(1, 3, 4),)
- verify_model(SymSizeInt(dim=1), example_args_1, {}, Expected1)
- verify_model(SymSizeInt(dim=-2), example_args_1, {}, Expected1)
-
- class SymSizeIntDynamic(Module):
- def __init__(self, dim):
- super().__init__()
- self.dim = dim
-
- def forward(self, x):
- shape_dim = torch.ops.aten.sym_size.int(x, self.dim)
+ shape_dim = torch.ops.aten.sym_size.int(x, 0)
return x.reshape(shape_dim, -1)
@I.ir_module
- class Expected2:
+ class Expected:
@R.function
def main(x: R.Tensor(("s0", 3, 4), dtype="float32")) -> R.Tuple(
R.Tensor(("s0", 12), dtype="float32")
@@ -8865,13 +8248,13 @@ def test_sym_size_int():
R.output(gv)
return gv
- example_args_2 = (torch.randn(2, 3, 4),)
+ example_args = (torch.randn(2, 3, 4),)
dynamic_shapes = {"x": {0: torch.export.Dim("dim")}}
verify_model(
- SymSizeIntDynamic(dim=0),
- example_args_2,
+ SymSizeInt(),
+ example_args,
{},
- Expected2,
+ Expected,
dynamic_shapes=dynamic_shapes,
map_free_vars=True,
)
diff --git a/tests/python/relax/test_frontend_from_fx.py
b/tests/python/relax/test_frontend_from_fx.py
index e68a4fa768..a489977958 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -14,8 +14,6 @@
# 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
@@ -33,7 +31,6 @@ 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):
@@ -901,14 +898,7 @@ def test_outer():
verify_model(Outer(), input_infos, {}, expected)
[email protected]
[email protected](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__()
@@ -939,14 +929,7 @@ def test_softplus():
verify_model(Softplus1(), input_info, {}, expected)
[email protected]
[email protected](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__()
@@ -1835,10 +1818,6 @@ def test_stochastic_depth():
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
@@ -1851,8 +1830,9 @@ def test_stochastic_depth():
R.output(gv)
return gv
+ # The PyTorch frontend imports models with inference semantics, so
stochastic
+ # depth is lowered to an identity even when the traced module is in
training mode.
verify_model(StochasticDepth1(), input_info, {}, expected1)
- verify_model(StochasticDepth2(), input_info, {}, expected1)
def test_layernorm():
@@ -2131,12 +2111,6 @@ def test_functional_cross_entropy():
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):
@@ -2180,9 +2154,6 @@ def test_groupnorm():
def test_instancenorm2d():
- torch.set_grad_enabled(False)
- torch.random.manual_seed(0)
-
input_info = [([1, 3, 10, 10], "float32")]
class InstanceNorm2d(Module):
@@ -2216,8 +2187,6 @@ def test_instancenorm2d():
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(),
@@ -2724,7 +2693,6 @@ operator_basic_unary = [
(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),
@@ -3116,15 +3084,6 @@ def test_extended_unary_ops():
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):
@@ -3771,41 +3730,6 @@ def test_interpolate():
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(
@@ -3911,44 +3835,6 @@ def test_interpolate():
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):
@@ -4401,65 +4287,28 @@ def test_masked_fill_inplace():
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"),
- )
-
[email protected](
+ "torch_dtype, expected_dtype",
+ [(torch.float32, "float32"), (None, "int64")],
+ ids=["float32", "default-int64"],
+)
+def test_get_attr_scalar_tensor_constant(torch_dtype, expected_dtype):
+ """Import rank-0 tensor constants folded to get_attr by FX tracing."""
-def test_empty():
- class Empty(Module):
+ class ScalarTensor(Module):
def forward(self, input):
- return torch.empty((10, 10), dtype=torch.float32)
+ if torch_dtype is None:
+ return torch.tensor(3)
+ return torch.tensor(3, dtype=torch_dtype)
- graph_model = fx.symbolic_trace(Empty())
+ graph_model = fx.symbolic_trace(ScalarTensor())
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"
+ bindings = mod["main"].body.blocks[0].bindings
+ assert len(bindings) == 1
+ value = bindings[0].value
+ assert isinstance(value, relax.Constant)
+ assert value.data.shape == ()
+ assert value.data.dtype == expected_dtype
def test_new_ones():
@@ -4513,10 +4362,6 @@ def test_new_zeros():
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)
@@ -4534,7 +4379,6 @@ def test_expand():
R.output(gv)
return gv
- verify_model(Expand1(), input_info, {}, expected1)
verify_model(Expand2(), input_info, {}, expected1)
@@ -4832,7 +4676,6 @@ def test_repeat():
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():
@@ -5696,18 +5539,6 @@ def test_gather():
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
@@ -5721,220 +5552,32 @@ def test_gather():
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")]
+ class IndexPut(Module):
+ def forward(self, data, indices, values):
+ return data.index_put_((indices,), values, accumulate=False)
@I.ir_module
- class Expected1D:
+ class Expected:
@R.function
def main(
data: R.Tensor((64,), dtype="float32"),
- indices_0: R.Tensor((128,), dtype="int64"),
+ indices: 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
+ data, R.tuple(indices), 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)
+ input_info = [((64,), "float32"), ((128,), "int64"), ((128,), "float32")]
+ verify_model(IndexPut(), input_info, {}, Expected)
def test_flip():
@@ -5942,10 +5585,6 @@ def test_flip():
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
@@ -5958,20 +5597,7 @@ def test_flip():
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():
@@ -6713,42 +6339,6 @@ def test_dtypes(torch_dtype, relax_dtype):
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")]
@@ -6806,7 +6396,7 @@ def test_round():
]
)
- for decimals in [0, 1, 2, 3]:
+ for decimals in [0, 2]:
torch_model = Round(decimals)
graph_model = fx.symbolic_trace(torch_model)
with torch.no_grad():
diff --git a/tests/python/relax/test_frontend_tflite.py
b/tests/python/relax/test_frontend_tflite.py
index e030f73b8a..1561826b17 100644
--- a/tests/python/relax/test_frontend_tflite.py
+++ b/tests/python/relax/test_frontend_tflite.py
@@ -1327,6 +1327,48 @@ def test_binary(tf_op, relax_op):
verify(Binary, Expected)
+def test_static_broadcast_lowered_to_multiply():
+ """Static TensorFlow broadcasts become TFLite MUL with a constant
operand."""
+
+ class RankExpansion(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2),
dtype=tf.float32)])
+ def func(self, x):
+ return tf.broadcast_to(x, [3, 2, 2])
+
+ @I.ir_module
+ class ExpectedRankExpansion:
+ @R.function
+ def main(x: R.Tensor((2, 2), dtype="float32")) -> R.Tensor((3, 2, 2),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((3, 2, 2), dtype="float32") = R.multiply(
+ x, R.const(np.ones((3, 2, 2), dtype="float32"))
+ )
+ R.output(gv)
+ return gv
+
+ verify(RankExpansion, ExpectedRankExpansion)
+
+ class ScalarInt(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.int32)])
+ def func(self, x):
+ return tf.broadcast_to(x, [4, 4])
+
+ @I.ir_module
+ class ExpectedScalarInt:
+ @R.function
+ def main(x: R.Tensor((), dtype="int32")) -> R.Tensor((4, 4),
dtype="int32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((4, 4), dtype="int32") = R.multiply(
+ x, R.const(np.ones((4, 4), dtype="int32"))
+ )
+ R.output(gv)
+ return gv
+
+ verify(ScalarInt, ExpectedScalarInt)
+
+
def test_pow():
class TfInput(tf.Module):
@tf.function(input_signature=[tf.TensorSpec(shape=(1, 30),
dtype=tf.float32)])
@@ -2056,6 +2098,33 @@ def test_gather():
verify(Gather, Expected)
+ # TensorFlow lowers embedding lookup to TFLite GATHER. Keep a case with
+ # constant params and multidimensional int32 indices, since the case above
+ # uses runtime params and one-dimensional int64 indices along axis 1.
+ class GatherConstantParams(tf.Module):
+ @tf.function(input_signature=[tf.TensorSpec(shape=(2, 3),
dtype=tf.int32)])
+ def func(self, indices):
+ params = tf.constant([[1, 2], [3, 4], [5, 6], [7, 8]],
dtype=tf.float32)
+ return tf.gather(params, indices, axis=0)
+
+ @I.ir_module
+ class ExpectedConstantParams:
+ @R.function
+ def main(indices: R.Tensor((2, 3), dtype="int32")) -> R.Tensor((2, 3,
2), dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ lv: R.Tensor((2, 3), dtype="int32") = R.astype(indices,
dtype="int32")
+ gv: R.Tensor((2, 3, 2), dtype="float32") = R.take(
+ R.const(np.array([[1, 2], [3, 4], [5, 6], [7, 8]],
dtype=np.float32)),
+ lv,
+ axis=0,
+ mode="fast",
+ )
+ R.output(gv)
+ return gv
+
+ verify(GatherConstantParams, ExpectedConstantParams)
+
def test_gather_nd():
class GatherND(tf.Module):
