Function _test_optimizer

tests/python/relax/test_training_optimizer_numeric.py:68–79

(target, dev, np_func, opt_type, *args, **kwargs)

@tvm.testing.parametrize_targets("llvm")
def _test_optimizer(target, dev, np_func, opt_type, *args, **kwargs):
    x = relax.Var("x", R.Tensor((3, 3), "float32"))
    y = relax.Var("y", R.Tensor((3,), "float32"))
    opt = opt_type(*args, **kwargs).init([x, y])
    mod = IRModule.from_expr(opt.get_function().with_attr("global_symbol", "main"))
    tvm_func = _legalize_and_build(mod, target, dev)["main"]

    param_arr = [np.random.rand(3, 3).astype(np.float32), np.random.rand(3).astype(np.float32)]
    grad_arr = [np.random.rand(3, 3).astype(np.float32), np.random.rand(3).astype(np.float32)]
    state_arr = _tvm_to_numpy(opt.state)

    _assert_run_result_same(tvm_func, np_func, [param_arr, grad_arr, state_arr])


lr, weight_decay = tvm.testing.parameters(