import torch import torch.nn as nn

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

==============================================================================

Multi-layer perceptron (MLP) models.

class Mlp1LayerModule(torch.nn.Module): def init(self): super().init() # Reset seed to make model deterministic. torch.manual_seed(0) self.fc0 = nn.Linear(3, 5) self.tanh0 = nn.Tanh() @export @annotate_args([ None, ([-1, -1], torch.float32, True), ]) def forward(self, x): return self.tanh0(self.fc0(x))

@register_test_case(module_factory=lambda: Mlp1LayerModule()) def Mlp1LayerModule_basic(module, tu: TestUtils): module.forward(tu.rand(5, 3))