Function test_invalid_mod

tests/python/relax/test_training_setup_trainer.py:183–229

()

def test_invalid_mod():
    @I.ir_module
    class NoAttr:
        @R.function
        def backbone(
            w0: R.Tensor((10, 5), "float32"),
            b0: R.Tensor((5,), "float32"),
            x: R.Tensor((1, 10), "float32"),
        ):
            with R.dataflow():
                lv0 = R.matmul(x, w0)
                gv = R.add(lv0, b0)
                out = R.nn.relu(gv)
                R.output(gv, out)
            return gv, out

    pred_sinfo = relax.TensorStructInfo((1, 5), "float32")
    setup_trainer = SetupTrainer(
        MSELoss(reduction="sum"),
        SGD(0.001),
        [pred_sinfo, pred_sinfo],
    )

    with pytest.raises((RuntimeError, ValueError)):
        SetupTrainer(
            MSELoss(reduction="sum"),
            SGD(0.001),
            [pred_sinfo, pred_sinfo],
        )(NoAttr)

    @I.ir_module
    class WrongFuncName:
        @R.function
        def main(
            w0: R.Tensor((10, 5), "float32"),
            b0: R.Tensor((5,), "float32"),
            x: R.Tensor((1, 10), "float32"),
        ):
            with R.dataflow():
                lv0 = R.matmul(x, w0)
                lv1 = R.add(lv0, b0)
                out = R.nn.relu(lv1)
                R.output(out)
            return out

    with pytest.raises(ValueError):
        setup_trainer(WrongFuncName)


if __name__ == "__main__":