Function test_adam_complex

tests/python/relax/test_training_optimizer.py:399–498

()

def test_adam_complex():
    x = relax.Var("x", R.Tensor((3, 3), "float32"))
    y = relax.Var("y", R.Tensor((3,), "float32"))
    adam = Adam(0.01, (0.8, 0.85), 1e-7, 0.1).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.8, "float32"))
            lv1: R.Tensor((), "float32") = optim_states[2]
            beta2_prod: R.Tensor((), "float32") = R.multiply(lv1, R.const(0.85, "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.1, "float32"), x)
            x_grad_new: R.Tensor((3, 3), "float32") = R.add(lv2, x_grad)
            lv3: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.8, "float32"), x_m)
            lv4: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.2, "float32"), x_grad_new)
            x_m_new: R.Tensor((3, 3), "float32") = R.add(lv3, lv4)
            lv5: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.85, "float32"), x_v)
            lv6: R.Tensor((3, 3), "float32") = R.multiply(x_grad_new, x_grad_new)
            lv7: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.15, "float32"), lv6)
            x_v_new: R.Tensor((3, 3), "float32") = R.add(lv5, lv7)
            lv8: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta1_prod)
            x_m_hat: R.Tensor((3, 3), "float32") = R.divide(x_m_new, lv8)
            lv9: R.Tensor((), "float32") = R.subtract(R.const(1, "float32"), beta2_prod)
            x_v_hat: R.Tensor((3, 3), "float32") = R.divide(x_v_new, lv9)
            lv10: R.Tensor((3, 3), "float32") = R.sqrt(x_v_hat)
            lv11: R.Tensor((3, 3), "float32") = R.add(lv10, R.const(1e-07, "float32"))