function_learning.py

ARITHMETIC_FUNCTIONS = {
    'add': lambda x, y: x + y,
    'sub': lambda x, y: x - y,
    'mul': lambda x, y: x * y,
    'div': lambda x, y: x / y,
    'squared': lambda x, y: torch.pow(x, 2),
    'root': lambda x, y: torch.sqrt(x),
}

for fn_str, fn in ARITHMETIC_FUNCTIONS.items():
       results[fn_str] = []

       # dataset
       X_train, y_train, X_test, y_test = generate_data(
            num_train=500, num_test=50,
            dim=100, num_sum=5, fn=fn,
            support=RANGE,
       )

       # models (Baseline models, NALU, NAC)
       for net in models:
           optim = torch.optim.RMSprop(net.parameters(), lr=LEARNING_RATE)
           train(net, optim, X_train, y_train, NUM_ITERS)
           mse = test(net, X_test, y_test).mean().item()
           results[fn_str].append(mse)