Function np_func

tests/python/relax/test_training_optimizer_numeric.py:90–98

(param_tuple, grad_tuple, state_tuple)

@tvm.testing.parametrize_targets("llvm")
def test_sgd(target, dev, lr, weight_decay):
    def np_func(param_tuple, grad_tuple, state_tuple):
        num_steps = state_tuple[0]
        param_tuple_new, state_tuple_new = [], []
        state_tuple_new.append(num_steps + 1)
        for i in range(len(param_tuple)):
            param = param_tuple[i]
            grad = grad_tuple[i]
            param_tuple_new.append(param - lr * (grad + weight_decay * param))
        return param_tuple_new, state_tuple_new

    _test_optimizer(target, dev, np_func, SGD, lr, weight_decay)