Function test_softplus

tests/python/relax/test_frontend_from_exported_program.py:885–929

()

def test_softplus():
    import torch
    from torch.nn import Module

    torch.set_grad_enabled(False)

    class Softplus0(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.softplus = torch.nn.Softplus(1.0, 20.0)

        def forward(self, x):
            return self.softplus(x)

    class Softplus1(Module):
        def forward(self, input):
            return torch.nn.functional.softplus(input, 1.0, 20.0)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(x: 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.multiply(
                    x, R.const(1.0, "float32")
                )
                lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(lv)
                lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv1, R.const(1.0, "float32"))
                lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(lv2)
                lv4: R.Tensor((1, 3, 10, 10), dtype="float32") = R.divide(
                    lv3, R.const(1.0, "float32")
                )
                lv5: R.Tensor((1, 3, 10, 10), dtype="bool") = R.greater(
                    lv, R.const(20.0, "float32")
                )
                lv6: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv5, x, lv4)
                gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv6,)
                R.output(gv)
            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)


def test_leakyrelu():