Function test_simple

tests/python/relax/test_training_setup_trainer.py:31–99

()

def test_simple():
    # fmt: off
    @I.ir_module
    class Backbone:
        I.module_attrs({"param_num": 1, "state_num": 0})
        @R.function
        def backbone(x: R.Tensor((2, 2), "float64"), y: R.Tensor((2, 2), "float64")):
            with R.dataflow():
                x1 = x + y
                R.output(x1)
            return x1

    @I.ir_module
    class Expected:
        I.module_attrs({"input_num": 1, "param_num": 1, "state_num": 0})
        @R.function
        def backbone(x: R.Tensor((2, 2), dtype="float64"), y: R.Tensor((2, 2), dtype="float64")) -> R.Tensor((2, 2), dtype="float64"):
            with R.dataflow():
                x1: R.Tensor((2, 2), dtype="float64") = R.add(x, y)
                R.output(x1)
            return x1

        @R.function
        def backbone_loss(x: R.Tensor((2, 2), dtype="float64"), y: R.Tensor((2, 2), dtype="float64"), targets: R.Tensor((2, 2), dtype="float64")) -> R.Tensor((), dtype="float64"):
            with R.dataflow():
                x1: R.Tensor((2, 2), dtype="float64") = R.add(x, y)
                lv: R.Tensor((2, 2), dtype="float64") = R.subtract(x1, targets)
                lv1: R.Tensor((2, 2), dtype="float64") = R.multiply(lv, lv)
                gv: R.Tensor((), dtype="float64") = R.sum(lv1, axis=None, keepdims=False)
                R.output(gv)
            return gv

        @R.function
        def backbone_loss_adjoint(x: R.Tensor((2, 2), dtype="float64"), y: R.Tensor((2, 2), dtype="float64"), targets: R.Tensor((2, 2), dtype="float64")) -> R.Tuple(R.Tensor((), dtype="float64"), R.Tuple(R.Tensor((2, 2), dtype="float64"))):
            with R.dataflow():
                x1: R.Tensor((2, 2), dtype="float64") = R.add(x, y)
                lv: R.Tensor((2, 2), dtype="float64") = R.subtract(x1, targets)
                lv1: R.Tensor((2, 2), dtype="float64") = R.multiply(lv, lv)
                gv: R.Tensor((), dtype="float64") = R.sum(lv1, axis=None, keepdims=False)
                gv_adjoint: R.Tensor((), dtype="float64") = R.ones(R.shape([]), dtype="float64")
                lv1_adjoint: R.Tensor((2, 2), dtype="float64") = R.broadcast_to(gv_adjoint, R.shape([2, 2]))
                lv_adjoint: R.Tensor((2, 2), dtype="float64") = R.multiply(lv1_adjoint, lv)
                lv_1: R.Tensor((2, 2), dtype="float64") = R.multiply(lv1_adjoint, lv)
                lv_adjoint1: R.Tensor((2, 2), dtype="float64") = R.add(lv_adjoint, lv_1)
                x1_adjoint: R.Tensor((2, 2), dtype="float64") = lv_adjoint1
                y_adjoint: R.Tensor((2, 2), dtype="float64") = x1_adjoint
                y_adjoint_out: R.Tensor((2, 2), dtype="float64") = y_adjoint
                R.output(gv, y_adjoint_out)
            return (gv, (y_adjoint_out,))

        @R.function
        def optimizer(params: R.Tuple(R.Tensor((2, 2), dtype="float64")), gradients: R.Tuple(R.Tensor((2, 2), dtype="float64")), optim_states: R.Tuple(R.Tensor((), dtype="int64"))) -> R.Tuple(R.Tuple(R.Tensor((2, 2), dtype="float64")), R.Tuple(R.Tensor((), dtype="int64"))):
            with R.dataflow():
                num_steps: R.Tensor((), dtype="int64") = optim_states[0]
                num_steps_new: R.Tensor((), dtype="int64") = R.add(num_steps, R.const(1, "int64"))
                y: R.Tensor((2, 2), dtype="float64") = params[0]
                y_grad: R.Tensor((2, 2), dtype="float64") = gradients[0]
                lv: R.Tensor((2, 2), dtype="float64") = R.multiply(R.const(0.10000000000000001, "float64"), y_grad)