Function test_logaddexp

tests/python/relax/test_frontend_from_exported_program.py:974–1016

()

def test_logaddexp():
    class LogAddExp(Module):
        def forward(self, input1, input2):
            return torch.logaddexp(input1, input2)

    @tvm.script.ir_module
    class expected:
        @R.function
        def main(
            input1: R.Tensor((1, 3, 10, 10), dtype="float32"),
            input2: R.Tensor((1, 3, 10, 10), dtype="float32"),
        ) -> R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")):
            # block 0
            with R.dataflow():
                lv: R.Tensor((1, 3, 10, 10), dtype="bool") = R.greater_equal(input1, input2)
                lv1: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv, input1, input2)
                lv2: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv, input2, input1)
                lv3: R.Tensor((1, 3, 10, 10), dtype="float32") = R.abs(input1)
                lv4: R.Tensor((1, 3, 10, 10), dtype="bool") = R.not_equal(
                    lv3, R.const(float("inf"), "float32")
                )
                lv5: R.Tensor((1, 3, 10, 10), dtype="bool") = R.equal(input1, input1)
                lv6: R.Tensor((1, 3, 10, 10), dtype="bool") = R.multiply(lv5, lv4)
                lv7: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_not(lv6)
                lv8: R.Tensor((1, 3, 10, 10), dtype="bool") = R.equal(input1, input2)
                lv9: R.Tensor((1, 3, 10, 10), dtype="bool") = R.logical_and(lv7, lv8)
                lv10: R.Tensor((1, 3, 10, 10), dtype="float32") = R.subtract(lv2, lv1)
                lv11: R.Tensor((1, 3, 10, 10), dtype="float32") = R.exp(lv10)
                lv12: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(
                    lv11, R.const(1.0, "float32")
                )
                lv13: R.Tensor((1, 3, 10, 10), dtype="float32") = R.log(lv12)
                lv14: R.Tensor((1, 3, 10, 10), dtype="float32") = R.add(lv1, lv13)
                lv15: R.Tensor((1, 3, 10, 10), dtype="float32") = R.where(lv9, input1, lv14)
                gv: R.Tuple(R.Tensor((1, 3, 10, 10), dtype="float32")) = (lv15,)
                R.output(gv)
            return gv

    example_args = (
        torch.randn(1, 3, 10, 10, dtype=torch.float32),
        torch.randn(1, 3, 10, 10, dtype=torch.float32),
    )
    verify_model(LogAddExp(), example_args, {}, expected)


def test_logical_and():