Function test_adam_simple

tests/python/relax/test_training_optimizer.py:301–396

()

def test_adam_simple():
    x = relax.Var("x", R.Tensor((3, 3), "float32"))
    y = relax.Var("y", R.Tensor((3,), "float32"))
    adam = Adam(0.01).init([x, y]).get_function()

    @R.function
    def adam_expected(
        params: R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
        gradients: R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
        optim_states: R.Tuple(
            R.Tensor((), "int64"),
            R.Tensor((), "float32"),
            R.Tensor((), "float32"),
            R.Tensor((3, 3), "float32"),
            R.Tensor((3,), "float32"),
            R.Tensor((3, 3), "float32"),
            R.Tensor((3,), "float32"),
        ),
    ) -> R.Tuple(
        R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
        R.Tuple(
            R.Tensor((), "int64"),
            R.Tensor((), "float32"),
            R.Tensor((), "float32"),
            R.Tensor((3, 3), "float32"),
            R.Tensor((3,), "float32"),
            R.Tensor((3, 3), "float32"),
            R.Tensor((3,), "float32"),
        ),
    ):
        R.func_attr({"global_symbol": "Adam"})
        # block 0
        with R.dataflow():
            num_steps: R.Tensor((), "int64") = optim_states[0]
            num_steps_new: R.Tensor((), "int64") = R.add(num_steps, R.const(1, "int64"))
            lv: R.Tensor((), "float32") = optim_states[1]
            beta1_prod: R.Tensor((), "float32") = R.multiply(lv, R.const(0.9, "float32"))
            lv1: R.Tensor((), "float32") = optim_states[2]
            beta2_prod: R.Tensor((), "float32") = R.multiply(lv1, R.const(0.999, "float32"))
            x: R.Tensor((3, 3), "float32") = params[0]
            x_grad: R.Tensor((3, 3), "float32") = gradients[0]
            x_m: R.Tensor((3, 3), "float32") = optim_states[3]
            x_v: R.Tensor((3, 3), "float32") = optim_states[5]
            lv2: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.9, "float32"), x_m)
            lv3: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.1, "float32"), x_grad)
            x_m_new: R.Tensor((3, 3), "float32") = R.add(lv2, lv3)
            lv4: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.999, "float32"), x_v)
            lv5: R.Tensor((3, 3), "float32") = R.multiply(x_grad, x_grad)
            lv6: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.001, "float32"), lv5)
            x_v_new: R.Tensor((3, 3), "float32") = R.add(lv4, lv6)
            lv7: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta1_prod)
            x_m_hat: R.Tensor((3, 3), "float32") = R.divide(x_m_new, lv7)
            lv8: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta2_prod)
            x_v_hat: R.Tensor((3, 3), "float32") = R.divide(x_v_new, lv8)
            lv9: R.Tensor((3, 3), "float32") = R.sqrt(x_v_hat)
            lv10: R.Tensor((3, 3), "float32") = R.add(lv9, R.const(1e-08, "float32"))
            lv11: R.Tensor((3, 3), "float32") = R.divide(x_m_hat, lv10)
            lv12: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.01, "float32"), lv11)