torch-mlir/e2e_testing/torchscript/basic.py

1310 lines
36 KiB
Python
Raw Normal View History

# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# Also available under a BSD-style license. See LICENSE.
import torch
from torch_mlir_e2e_test.torchscript.framework import TestUtils
from torch_mlir_e2e_test.torchscript.registry import register_test_case
from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
# ==============================================================================
class MmModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
Introduce `!torch.tensor` / `!torch.vtensor` types. This removes our reliance on the numpy dialect and avoids our off-label use of the builtin tnesor type for modeling unknown dtypes. The `!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor. The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic tensor. The new types look as follows syntactically: ``` // Least-static-information, non-value-semantic tensor. !torch.tensor // Explicit form of least-static-information variant. !torch.tensor<*,unk> // Least-static-information, value-semantic tensor. !torch.vtensor // Explicit form of least-static-information variant. !torch.vtensor<*,unk> // Fixed-set of allowable element types, with first-class support for // Torch's frontend signedness semantics. !torch.tensor<*,si32> // First-class support for unknown dtypes. !torch.tensor<[?,?,?],unk> // Standard MLIR representation of `?` for unknown dimensions. !torch.tensor<[?,2,?,4],unk> // Statically shaped / dtyped example. !torch.vtensor<[1,2,3,4],f32> ``` This required fairly significant changes throughout the compiler, but overall it is a big cleanup. We now have a much clearer layering of "the Torch frontend lowering" vs "lowering to std + linalg + etc.". At the C++ level, there is `ValueTensorType`, `NonValueTensorType`. We also have a helper `BaseTensorType` (kind of like ShapedType) which interoperates with those two. Included changes: - New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for creating torch tensor literals in the frontend. - Consistently use signedness for the types (except i1 which I didn't touch -- we need to sort out the situation with !basicpy.BoolType there anyway so will be attending to that soon) - Frontend can annotate whether an argument to the function has value semantics. We currently require this, as our backend contract does not currently allow us to even model the non-value-semantic case. Before, the value-semantic assumption was randomly injected in the middle of the pass pipeline. - Move ArrayToTensor (now called MaximizeValueSemantics) and RefinePublicReturn passes to torch dialect. - The TorchToStd and TorchToLinalg passes are now type conversions from `!torch.vtensor` to `tensor` and use the dialect conversion infra. The overall conversion pipeline is set up following the best practices of the "Type Conversions the Not-So-Hard Way" talk. This required introducing `torch-func-builtin-tensorize` and `torch-finalizing-builtin-tensorize` passes analogous to the upstream bufferization passes with the corresponding names (mostly just copypasta from there). - Misc Torch-level canonicalizations -- we now cleanly layer the lowering to std later in the pipeline, so we are gradually lessening our reliance on random std constant folding before we get to that point. Recommended review order: - New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp - New ops in TorchOps.td / TorchOps.cpp - Less important / more mechanical stuff - Frontend changes. - Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
])
def forward(self, lhs, rhs):
return torch.mm(lhs, rhs)
@register_test_case(module_factory=lambda: MmModule())
def MmModule_basic(module, tu: TestUtils):
module.forward(tu.rand(4, 4), tu.rand(4, 4))
2021-09-24 05:50:37 +08:00
@register_test_case(module_factory=lambda: MmModule())
def MmModule_chained(module, tu: TestUtils):
res = module.forward(tu.rand(4, 4), tu.rand(4, 4))
module.forward(res, res)
# ==============================================================================
class BmmModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
([-1, -1, -1], torch.float32, True),
])
def forward(self, lhs, rhs):
return torch.bmm(lhs, rhs)
@register_test_case(module_factory=lambda: BmmModule())
def BmmModule_basic(module, tu: TestUtils):
module.forward(tu.rand(3, 4, 5), tu.rand(3, 5, 4))
# ==============================================================================
# A subgraph with multiple mm ops.
class MmDagModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([4, 4], torch.float32, True),
([4, 4], torch.float32, True),
])
def forward(self, lhs, rhs):
return torch.mm(lhs, torch.mm(lhs, rhs))
@register_test_case(module_factory=lambda: MmDagModule())
def MmDagModule_basic(module, tu: TestUtils):
module.forward(tu.rand(4, 4), tu.rand(4, 4))
# ==============================================================================
class MmTanhModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
Introduce `!torch.tensor` / `!torch.vtensor` types. This removes our reliance on the numpy dialect and avoids our off-label use of the builtin tnesor type for modeling unknown dtypes. The `!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor. The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic tensor. The new types look as follows syntactically: ``` // Least-static-information, non-value-semantic tensor. !torch.tensor // Explicit form of least-static-information variant. !torch.tensor<*,unk> // Least-static-information, value-semantic tensor. !torch.vtensor // Explicit form of least-static-information variant. !torch.vtensor<*,unk> // Fixed-set of allowable element types, with first-class support for // Torch's frontend signedness semantics. !torch.tensor<*,si32> // First-class support for unknown dtypes. !torch.tensor<[?,?,?],unk> // Standard MLIR representation of `?` for unknown dimensions. !torch.tensor<[?,2,?,4],unk> // Statically shaped / dtyped example. !torch.vtensor<[1,2,3,4],f32> ``` This required fairly significant changes throughout the compiler, but overall it is a big cleanup. We now have a much clearer layering of "the Torch frontend lowering" vs "lowering to std + linalg + etc.". At the C++ level, there is `ValueTensorType`, `NonValueTensorType`. We also have a helper `BaseTensorType` (kind of like ShapedType) which interoperates with those two. Included changes: - New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for creating torch tensor literals in the frontend. - Consistently use signedness for the types (except i1 which I didn't touch -- we need to sort out the situation with !basicpy.BoolType there anyway so will be attending to that soon) - Frontend can annotate whether an argument to the function has value semantics. We currently require this, as our backend contract does not currently allow us to even model the non-value-semantic case. Before, the value-semantic assumption was randomly injected in the middle of the pass pipeline. - Move ArrayToTensor (now called MaximizeValueSemantics) and RefinePublicReturn passes to torch dialect. - The TorchToStd and TorchToLinalg passes are now type conversions from `!torch.vtensor` to `tensor` and use the dialect conversion infra. The overall conversion pipeline is set up following the best practices of the "Type Conversions the Not-So-Hard Way" talk. This required introducing `torch-func-builtin-tensorize` and `torch-finalizing-builtin-tensorize` passes analogous to the upstream bufferization passes with the corresponding names (mostly just copypasta from there). - Misc Torch-level canonicalizations -- we now cleanly layer the lowering to std later in the pipeline, so we are gradually lessening our reliance on random std constant folding before we get to that point. Recommended review order: - New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp - New ops in TorchOps.td / TorchOps.cpp - Less important / more mechanical stuff - Frontend changes. - Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
])
def forward(self, lhs, rhs):
return torch.tanh(self.matmul(lhs, rhs))
def matmul(self, lhs, rhs):
return torch.mm(lhs, rhs)
@register_test_case(module_factory=lambda: MmTanhModule())
def MmTanhModule_basic(module, tu: TestUtils):
module.forward(tu.rand(4, 2), tu.rand(2, 4))
# ==============================================================================
class AddmmModuleFloat(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
])
def forward(self, M, mat1, mat2):
return torch.addmm(M, mat1, mat2, beta=3.0, alpha=7.0)
@register_test_case(module_factory=lambda: AddmmModuleFloat())
def AddmmModuleFloat_basic(module, tu: TestUtils):
module.forward(tu.rand(4, 4), tu.rand(4, 2), tu.rand(2, 4))
# ==============================================================================
class AddmmModuleBroadcastable(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
])
def forward(self, M, mat1, mat2):
return torch.addmm(M, mat1, mat2, beta=2.0, alpha=7.0)
@register_test_case(module_factory=lambda: AddmmModuleBroadcastable())
def AddmmModule_broadcastable(module, tu: TestUtils):
module.forward(tu.rand(1, 2), tu.rand(3, 2), tu.rand(2, 2))
# ==============================================================================
class AddmmModuleDifferentRankBroadcastable(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1], torch.float32, True),
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
])
def forward(self, M, mat1, mat2):
return torch.addmm(M, mat1, mat2, beta=11.0, alpha=7.0)
@register_test_case(module_factory=lambda: AddmmModuleDifferentRankBroadcastable())
def AddmmModule_differentRankBroadcastable(module, tu: TestUtils):
module.forward(tu.rand(3), tu.rand(3, 2), tu.rand(2, 3))
# ==============================================================================
class AdaptiveAvgPool2dModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.aap2d = torch.nn.AdaptiveAvgPool2d((1, 1))
@export
@annotate_args([
None,
([-1, -1, -1, -1], torch.float32, True),
])
def forward(self, x):
return self.aap2d(x)
@register_test_case(module_factory=lambda: AdaptiveAvgPool2dModule())
def AdaptiveAvgPool2dModule_basic(module, tu: TestUtils):
module.forward(tu.rand(10, 3, 8, 9))
# ==============================================================================
class FlattenStaticModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.flat = torch.nn.Flatten(2, 4)
@export
@annotate_args([
None,
([10, 3, 8, 9, 3, 4], torch.float32, True),
])
def forward(self, x):
return self.flat(x)
@register_test_case(module_factory=lambda: FlattenStaticModule())
def FlattenStaticModule_basic(module, tu: TestUtils):
module.forward(tu.rand(10, 3, 8, 9, 3, 4))
# ==============================================================================
class FlattenRank0Module(torch.nn.Module):
def __init__(self):
super().__init__()
self.flat = torch.nn.Flatten(-1, -1)
@export
@annotate_args([
None,
([], torch.float32, True),
])
def forward(self, x):
return self.flat(x)
@register_test_case(module_factory=lambda: FlattenRank0Module())
def FlattenRank0Module_basic(module, tu: TestUtils):
module.forward(torch.tensor(4.0))
# ==============================================================================
class FlattenDynamicModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.flat = torch.nn.Flatten(2, 4)
@export
@annotate_args([
None,
([-1, -1, -1, 9, 3, -1], torch.float32, True),
])
def forward(self, x):
return self.flat(x)
@register_test_case(module_factory=lambda: FlattenDynamicModule())
def FlattenDynamicModule_basic(module, tu: TestUtils):
module.forward(tu.rand(10, 3, 8, 9, 3, 4))
# ==============================================================================
class MaxPool2dModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.mp2d = torch.nn.MaxPool2d(kernel_size=[6, 8],
stride=[2, 2],
padding=[3, 4],
dilation=2)
@export
@annotate_args([
None,
([-1, -1, -1, -1], torch.float32, True),
])
def forward(self, x):
return self.mp2d(x)
# ==============================================================================
@register_test_case(module_factory=lambda: MaxPool2dModule())
def MaxPool2dModule_basic(module, tu: TestUtils):
module.forward(tu.rand(1, 1, 20, 20) - 0.5)
2021-09-17 14:49:04 +08:00
class MaxPool2dStaticModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.mp2d = torch.nn.MaxPool2d(kernel_size=[3, 3],
stride=[2, 2],
padding=[1, 1],
dilation=[1, 1])
@export
@annotate_args([
None,
([1, 64, 112, 112], torch.float32, True),
])
def forward(self, x):
return self.mp2d(x)
@register_test_case(module_factory=lambda: MaxPool2dStaticModule())
def MaxPool2dStaticModule_basic(module, tu: TestUtils):
module.forward(tu.rand(1, 64, 112, 112))
# ==============================================================================
class ConstantPad2dStaticModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.pad2d = torch.nn.ConstantPad2d((0, 1, 2, 3), -float('inf'))
@export
@annotate_args([
None,
([1, 1, 20, 20], torch.float32, True),
])
def forward(self, x):
return self.pad2d(x)
@register_test_case(module_factory=lambda: ConstantPad2dStaticModule())
def ConstantPad2dStaticModule_basic(module, tu: TestUtils):
module.forward(tu.rand(1, 1, 20, 20) - 0.5)
# ==============================================================================
class ConstantPadNdModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1, -1, -1, -1], torch.float32, True),
])
def forward(self, x):
return torch.ops.aten.constant_pad_nd(x, (0, 1), -float('inf'))
@register_test_case(module_factory=lambda: ConstantPadNdModule())
def ConstantPadNdModule_basic(module, tu: TestUtils):
module.forward(tu.rand(1, 1, 20, 20, 4, 4) - 0.5)
# ==============================================================================
class ConstantPadNdStaticModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([1, 1, 20, 20, 4, 4], torch.float32, True),
])
def forward(self, x):
return torch.ops.aten.constant_pad_nd(x, (0, 1), -float('inf'))
@register_test_case(module_factory=lambda: ConstantPadNdStaticModule())
def ConstantPadNdStaticModule_basic(module, tu: TestUtils):
module.forward(tu.rand(1, 1, 20, 20, 4, 4) - 0.5)
# ==============================================================================
class ConstantPadNdPartialStaticModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([1, 1, 20, 20, -1, -1], torch.float32, True),
])
def forward(self, x):
return torch.ops.aten.constant_pad_nd(x, (0, 1, 2, 3), -float('inf'))
@register_test_case(module_factory=lambda: ConstantPadNdPartialStaticModule())
def ConstantPadNdPartialStaticModule_basic(module, tu: TestUtils):
module.forward(tu.rand(1, 1, 20, 20, 4, 4) - 0.5)
# ==============================================================================
2021-09-17 14:49:04 +08:00
class TransposeIntModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([3, 4, 2], torch.float32, True),
])
def forward(self, x):
return torch.transpose(x, 0, 1)
@register_test_case(module_factory=lambda: TransposeIntModule())
def TransposeIntModule_basic(module, tu: TestUtils):
module.forward(tu.rand(3, 4, 2))
# ==============================================================================
class PermuteModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([3, 4, 2], torch.float32, True)
])
def forward(self, x):
return x.permute(0, 2, 1)
@register_test_case(module_factory=lambda: PermuteModule())
def PermuteModule_basic(module, tu: TestUtils):
module.forward(tu.rand(3, 4, 2))
# ==============================================================================
class TransposeIntNegDimsModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([3, 4, 2], torch.float32, True),
])
def forward(self, x):
return torch.transpose(x, -1, -2)
@register_test_case(module_factory=lambda: TransposeIntNegDimsModule())
def TransposeIntNegDimsModule_basic(module, tu: TestUtils):
module.forward(tu.rand(3, 4, 2))
# ==============================================================================
class PermuteNegativeIndexModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([3, 4, 2], torch.float32, True)
])
def forward(self, x):
return x.permute(0, -1, 1)
@register_test_case(module_factory=lambda: PermuteNegativeIndexModule())
def PermuteNegativeIndexModule_basic(module, tu: TestUtils):
module.forward(tu.rand(3, 4, 2))
# ==============================================================================
class TensorsConcatModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
([-1, -1, -1], torch.float32, True),
([-1, -1, -1], torch.float32, True),
])
def forward(self, x, y, z):
return torch.cat([x, y, z], 1)
@register_test_case(module_factory=lambda: TensorsConcatModule())
def TensorsConcatModule_basic(module, tu: TestUtils):
module.forward(tu.rand(2, 2, 4), tu.rand(2, 1, 4), tu.rand(2, 3, 4))
# ==============================================================================
class GatherModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
([-1, -1, -1], torch.int64, True),
])
def forward(self, tensor, indices):
return torch.gather(tensor, 2, indices)
@register_test_case(module_factory=lambda: GatherModule())
def GatherModule_basic(module, tu: TestUtils):
module.forward(tu.rand(2, 3, 4), torch.tensor([[[1, 2, 3], [1, 2, 3]]]))
# ==============================================================================
class GatherStaticModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([2, 3, 4], torch.float32, True),
([1, 2, 3], torch.int64, True),
])
def forward(self, tensor, indices):
return torch.gather(tensor, 2, indices)
@register_test_case(module_factory=lambda: GatherStaticModule())
def GatherStaticModule_basic(module, tu: TestUtils):
module.forward(tu.rand(2, 3, 4), torch.tensor([[[1, 2, 3], [1, 2, 3]]]))
# ==============================================================================
class AddSizeIntModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
])
def forward(self, tensor):
# This is a workaround for not supporting scalar arguments.
# TODO: pass in dim as an argument to the forward method when scalar
# arguments are supported.
return tensor.add(tensor, alpha=tensor.size(1))
@register_test_case(module_factory=lambda: AddSizeIntModule())
def AddSizeIntModule_basic(module, tu: TestUtils):
module.forward(torch.randn(3, 3))
# ==============================================================================
class AddSizeIntNegDimModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
])
def forward(self, tensor):
# This is a workaround for not supporting scalar arguments.
# TODO: pass in dim as an argument to the forward method when scalar
# arguments are supported.
return tensor.add(tensor, alpha=tensor.size(-2))
@register_test_case(module_factory=lambda: AddSizeIntNegDimModule())
def AddSizeIntNegDimModule_basic(module, tu: TestUtils):
module.forward(torch.randn(3, 3))
# ==============================================================================
class EmbeddingModule(torch.nn.Module):
def __init__(self):
super().__init__()
torch.manual_seed(0)
self.embed = torch.nn.Embedding(num_embeddings=100,
embedding_dim=50,
padding_idx=4)
@export
@annotate_args([
None,
([-1, -1], torch.int64, True),
])
def forward(self, indices):
return self.embed.forward(indices)
@register_test_case(module_factory=lambda: EmbeddingModule())
def EmbeddingModule_basic(module, tu: TestUtils):
module.forward(torch.randint(100, (3, 3)))
# ==============================================================================
class SoftmaxIntModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.softmax = torch.nn.Softmax(2)
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
])
def forward(self, tensor):
return self.softmax.forward(tensor)
@register_test_case(module_factory=lambda: SoftmaxIntModule())
def SoftmaxIntModule_basic(module, tu: TestUtils):
module.forward(torch.randn(3, 2, 4))
# ==============================================================================
class _SoftmaxModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
])
def forward(self, tensor):
return torch.ops.aten._softmax(tensor, 0, False)
@register_test_case(module_factory=lambda: _SoftmaxModule())
def _SoftmaxModule_basic(module, tu: TestUtils):
module.forward(torch.randn(3, 2, 4))
# ==============================================================================
class SoftmaxIntNegDimModule(torch.nn.Module):
def __init__(self):
super().__init__()
torch.manual_seed(0)
self.softmax = torch.nn.Softmax(-2)
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
])
def forward(self, tensor):
return self.softmax.forward(tensor)
@register_test_case(module_factory=lambda: SoftmaxIntNegDimModule())
def SoftmaxIntNegDimModule_basic(module, tu: TestUtils):
module.forward(torch.randn(3, 2, 4))
# ==============================================================================
class SoftmaxIntArgTypeF64Module(torch.nn.Module):
def __init__(self):
super().__init__()
torch.manual_seed(0)
self.softmax = torch.nn.Softmax(2)
@export
@annotate_args([
None,
([-1, -1, -1], torch.float64, True),
])
def forward(self, tensor):
return self.softmax.forward(tensor)
@register_test_case(module_factory=lambda: SoftmaxIntArgTypeF64Module())
def SoftmaxIntArgTypeF64Module_basic(module, tu: TestUtils):
module.forward(torch.randn(3, 2, 4).double())
2021-10-19 16:25:08 +08:00
# ==============================================================================
class _LogSoftmaxModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
])
def forward(self, tensor):
return torch.ops.aten._log_softmax(tensor, dim=0, half_to_float=False)
@register_test_case(module_factory=lambda: _LogSoftmaxModule())
def _LogSoftmaxModule_basic(module, tu: TestUtils):
module.forward(torch.randn(3, 2, 4))
class _LogSoftmaxModuleStable(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1], torch.float32, True),
])
def forward(self, tensor):
return torch.ops.aten._log_softmax(tensor, dim=0, half_to_float=False)
@register_test_case(module_factory=lambda: _LogSoftmaxModuleStable())
def _LogSoftmaxModuleStable_basic(module, tu: TestUtils):
# testing for numerical stability.
# Should result in tensor([-1e9, 0.00]) rather than tensor([-inf, 0.]).
a = torch.tensor([0, 1e9])
module.forward(a)
# ==============================================================================
class HardsigmoidModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
])
def forward(self, x):
return torch.ops.aten.hardsigmoid(x)
@register_test_case(module_factory=lambda: HardsigmoidModule())
def HardsigmoidModule_basic(module, tu: TestUtils):
module.forward(torch.tensor([[4.0, -5.0, 3.0], [2.9, -1.5, -3.0]]))
# ==============================================================================
class HardsigmoidRandomModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
])
def forward(self, x):
return torch.ops.aten.hardsigmoid(x)
@register_test_case(module_factory=lambda: HardsigmoidRandomModule())
def HardsigmoidRandomModule_basic(module, tu: TestUtils):
module.forward(tu.rand(3, 4, low=-10, high=10))
# ==============================================================================
2021-10-19 16:25:08 +08:00
class BroadcastToModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, 1], torch.float32, True),
])
def forward(self, x):
return torch.broadcast_to(x, [1, -1, -1, 4])
@register_test_case(module_factory=lambda: BroadcastToModule())
def BroadcastToModule_basic(module, tu: TestUtils):
2021-10-21 12:52:53 +08:00
module.forward(tu.rand(3, 1, 1))
# ==============================================================================
class ExpandModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, 1], torch.float32, True),
])
def forward(self, x):
return x.expand([1, -1, -1, 4])
@register_test_case(module_factory=lambda: ExpandModule())
def ExpandModule_basic(module, tu: TestUtils):
module.forward(tu.rand(3, 1, 1))
2021-10-29 15:15:05 +08:00
# ==============================================================================
2021-10-29 15:15:05 +08:00
class ContiguousModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
])
def forward(self, x):
return x.contiguous()
@register_test_case(module_factory=lambda: ContiguousModule())
def ContiguousModule_basic(module, tu: TestUtils):
module.forward(tu.rand(3, 1))
# ==============================================================================
class LogSoftmaxIntModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.log_softmax = torch.nn.LogSoftmax(2)
@export
@annotate_args([
None,
([-1, -1, -1], torch.float64, True),
])
def forward(self, tensor):
return self.log_softmax.forward(tensor)
@register_test_case(module_factory=lambda: LogSoftmaxIntModule())
def LogSoftmaxIntModule_basic(module, tu: TestUtils):
module.forward(torch.randn(3, 2, 4).double())
2021-11-09 04:28:51 +08:00
# ==============================================================================
class NumToTensorIntModule(torch.nn.Module):
2021-11-09 04:28:51 +08:00
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
])
def forward(self):
return torch.ops.prim.NumToTensor(1)
@register_test_case(module_factory=lambda: NumToTensorIntModule())
def NumToTensorIntModule_basic(module, tu: TestUtils):
module.forward()
# ==============================================================================
class NumToTensorFloatModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
])
def forward(self):
return torch.ops.prim.NumToTensor(1.0)
@register_test_case(module_factory=lambda: NumToTensorFloatModule())
def NumToTensorFloatModule_basic(module, tu: TestUtils):
2021-11-09 04:28:51 +08:00
module.forward()
# ==============================================================================
# This test can be removed once we have one real op returning 3 float32 tensors
class ReturnThreeTensorFloat32(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
])
def forward(self, a, b, c):
return a, b, c
@register_test_case(module_factory=lambda: ReturnThreeTensorFloat32())
def ReturnThreeTensorFloat32_basic(module, tu: TestUtils):
module.forward(tu.rand(2, 3), tu.rand(2, 3), tu.rand(2, 3))
# ==============================================================================
2021-11-25 06:01:48 +08:00
class AddCMulModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
2021-11-25 06:01:48 +08:00
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
])
def forward(self, input, tensor1, tensor2):
return torch.addcmul(input, tensor1, tensor2, value=1.0)
@register_test_case(module_factory=lambda: AddCMulModule())
def AddCMulModule_basic(module, tu: TestUtils):
module.forward(tu.rand(1,3), tu.rand(1,3), tu.rand(1,3))
# ==============================================================================
2021-11-25 06:01:48 +08:00
class AddCDivModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
2021-11-25 06:01:48 +08:00
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
([-1, -1], torch.float32, True),
])
def forward(self, input, tensor1, tensor2):
return torch.addcdiv(input, tensor1, tensor2, value=1.0)
@register_test_case(module_factory=lambda: AddCDivModule())
def AddCDivModule_basic(module, tu: TestUtils):
module.forward(tu.rand(1,3), tu.rand(1,3), tu.rand(1,3))
# ==============================================================================
class tensorIntModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
])
def forward(self):
a = 1
return torch.tensor(a)
@register_test_case(module_factory=lambda: tensorIntModule())
def TensorIntModule_basic(module, tu: TestUtils):
module.forward()
# ==============================================================================
class tensorFloatModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
])
def forward(self):
a = 1.0
return torch.tensor(a)
@register_test_case(module_factory=lambda: tensorFloatModule())
def TensorFloatModule_basic(module, tu: TestUtils):
module.forward()
# ==============================================================================
class DropoutModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
])
def forward(self, x):
return torch.dropout(x, 0.0, False)
@register_test_case(module_factory=lambda: DropoutModule())
def DropoutModule_basic(module, tu: TestUtils):
module.forward(tu.rand(3, 4))
# ==============================================================================
class MeanModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([3, 4], torch.float32, True),
])
def forward(self, x):
return torch.mean(x)
@register_test_case(module_factory=lambda: MeanModule())
def MeanModule_basic(module, tu: TestUtils):
module.forward(torch.randn(3, 4))
# ==============================================================================
class MeanDynamicSizesModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
])
def forward(self, x):
return torch.mean(x)
@register_test_case(module_factory=lambda: MeanDynamicSizesModule())
def MeanDynamicSizesModule_basic(module, tu: TestUtils):
module.forward(torch.randn(3, 4))
# ==============================================================================
class NumelModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
])
def forward(self, input):
return torch.numel(input)
@register_test_case(module_factory=lambda: NumelModule())
def NumelModule_basic(module, tu: TestUtils):
module.forward(tu.rand(4, 3, 5))
# ==============================================================================
class NumelZeroRankModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([], torch.int64, True),
])
def forward(self, input):
return torch.numel(input)
@register_test_case(module_factory=lambda: NumelZeroRankModule())
def NumelZeroRankModule_basic(module, tu: TestUtils):
module.forward(torch.randint(10,[]))
# ==============================================================================
class BoolTensorReturnFalseModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1], torch.bool, True),
])
def forward(self, a):
return a
@register_test_case(module_factory=lambda: BoolTensorReturnFalseModule())
def BoolTensorReturnFalseModule_basic(module, tu: TestUtils):
module.forward(torch.tensor([0, 0], dtype=torch.bool))
# ==============================================================================
class BoolTensorReturnTrueModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1], torch.bool, True),
])
def forward(self, a):
return a
@register_test_case(module_factory=lambda: BoolTensorReturnTrueModule())
def BoolTensorReturnTrueModule_basic(module, tu: TestUtils):
module.forward(torch.tensor([1, 1, 1, 1, 1], dtype=torch.bool))
# ==============================================================================
class BoolTensorReturnMixedModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.bool, True),
])
def forward(self, a):
return a
@register_test_case(module_factory=lambda: BoolTensorReturnMixedModule())
def BoolTensorReturnMixedModule_basic(module, tu: TestUtils):
module.forward(torch.tensor([[1, 0], [0,1]], dtype=torch.bool))
2021-12-17 12:08:07 +08:00
# ==============================================================================
2021-12-17 12:08:07 +08:00
class TModuleRank2(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
])
def forward(self, lhs):
return torch.t(lhs)
@register_test_case(module_factory=lambda: TModuleRank2())
def TModuleRank2_basic(module, tu: TestUtils):
module.forward(tu.rand(3, 4))
2022-01-30 01:10:50 +08:00
# ==============================================================================
2021-12-17 12:08:07 +08:00
class TModuleRank1(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1], torch.float32, True),
])
def forward(self, lhs):
return torch.t(lhs)
@register_test_case(module_factory=lambda: TModuleRank1())
def TModuleRank1_basic(module, tu: TestUtils):
module.forward(tu.rand(3))
2022-01-30 01:10:50 +08:00
# ==============================================================================
2021-12-17 12:08:07 +08:00
class TModuleRank0(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([], torch.float32, True),
])
def forward(self, lhs):
return torch.t(lhs)
@register_test_case(module_factory=lambda: TModuleRank0())
def TModuleRank0_basic(module, tu: TestUtils):
module.forward(torch.tensor(7, dtype=torch.float32))
2022-01-30 01:10:50 +08:00
# ==============================================================================
class TensorLiteralModule(torch.nn.Module):
def __init__(self):
super().__init__()
torch.manual_seed(0)
self.t = torch.randint(-5, 5, (2, 3))
@export
@annotate_args([
None,
])
def forward(self):
return torch.add(self.t, self.t)
@register_test_case(module_factory=lambda: TensorLiteralModule())
def TensorLiteralModule_basic(module, tu: TestUtils):
module.forward()
2022-01-30 01:10:50 +08:00
# ==============================================================================
class TensorOpaqueLiteralModule(torch.nn.Module):
def __init__(self):
super().__init__()
torch.manual_seed(0)
self.t = torch.randint(-5, 5, (256, 1024))
@export
@annotate_args([
None,
])
def forward(self):
return torch.add(self.t, self.t)
@register_test_case(module_factory=lambda: TensorOpaqueLiteralModule())
def TensorOpaqueLiteralModule_basic(module, tu: TestUtils):
module.forward()
2022-01-30 01:10:50 +08:00
# ==============================================================================
class ReturnTwoTensorF32I64(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1], torch.float32, True),
([-1, -1], torch.int64, True),
])
def forward(self, a, b):
return a, b
@register_test_case(module_factory=lambda: ReturnTwoTensorF32I64())
def ReturnTwoTensorF32I64_basic(module, tu: TestUtils):
module.forward(tu.rand(2, 3), torch.randint(5, (2, 3)))
2022-01-30 01:10:50 +08:00
# ==============================================================================
class IndexTensorModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1], torch.float32, True),
([-1, -1], torch.int64, True),
])
def forward(self, x, index):
return torch.ops.aten.index(x, (index,))
@register_test_case(module_factory=lambda: IndexTensorModule())
def IndexTensorModule_basic(module, tu: TestUtils):
module.forward(tu.rand(5), torch.randint(4, (2, 3)))
2022-01-30 01:10:50 +08:00
# ==============================================================================
class SquareModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
])
def forward(self, x):
return torch.ops.aten.square(x)
@register_test_case(module_factory=lambda: SquareModule())
def SquareModule_basic(module, tu: TestUtils):
module.forward(tu.rand(2, 3, 4))
# ==============================================================================
class VarUnbiasedModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
])
def forward(self, x):
return torch.ops.aten.var(x, unbiased=True)
@register_test_case(module_factory=lambda: VarUnbiasedModule())
def VarUnbiasedModule_basic(module, tu: TestUtils):
module.forward(tu.rand(2, 3, 4))
# ==============================================================================
class VarBiasedModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
])
def forward(self, x):
return torch.ops.aten.var(x, unbiased=False)
@register_test_case(module_factory=lambda: VarBiasedModule())
def VarBiasedModule_basic(module, tu: TestUtils):
module.forward(tu.rand(2, 3, 4))
# ==============================================================================
class StdUnbiasedModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
])
def forward(self, x):
return torch.ops.aten.std(x, unbiased=True)
@register_test_case(module_factory=lambda: StdUnbiasedModule())
def StdUnbiasedModule_basic(module, tu: TestUtils):
module.forward(tu.rand(2, 3, 4))
# ==============================================================================
class StdBiasedModule(torch.nn.Module):
def __init__(self):
super().__init__()
@export
@annotate_args([
None,
([-1, -1, -1], torch.float32, True),
])
def forward(self, x):
return torch.ops.aten.std(x, unbiased=False)
@register_test_case(module_factory=lambda: StdBiasedModule())
def StdBiasedModule_basic(module, tu: TestUtils):
module.forward(tu.rand(2, 3, 4))