Function test_hardtanh

tests/python/relax/test_frontend_from_exported_program.py:847–882

()

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)

    class Hardtanh3(torch.nn.Module):
        def forward(self, input):
            return torch.ops.aten.hardtanh_(input)

    @tvm.script.ir_module
    class expected_for_1_2:
        @R.function
        def main(inp_0: R.Tensor((1, 3, 10, 10), dtype="float32")) -> R.Tuple(
            R.Tensor((1, 3, 10, 10), dtype="float32")
        ):
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="float32") = R.clip(
                    inp_0, R.prim_value(T.float64(-1.0)), R.prim_value(T.float64(1.0))
                )
                gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv,)
                R.output(gv)
            return gv

    example_args = (torch.randn(1, 3, 10, 10, dtype=torch.float32),)
    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():