<jupyter_start><jupyter_text>Toy Model on Snelson Dataset<jupyter_code>df = pd.read_csv('snelson.csv')
train_x = torch.from_numpy(df.x.values[:, np.newaxis]).float().to(device)
train_y = torch.from_numpy(df.y.values).float().to(device)
train_x.shape, train_y.shape
def train(x, y, model, mll, optim, lanc_iter=100, pre_size=100):
  model.train()

  optim.zero_grad()

  with gp.settings.cg_tolerance(1e-2), \
       gp.settings.max_preconditioner_size(pre_size), \
       gp.settings.max_root_decomposition_size(lanc_iter):
    output = model(x)
    loss = -mll(output, y)

    loss.backward()

    for k, p in enumerate(model.parameters()):
        print(f'[{k}] [{p}] --> {p.grad}')

    optim.step()

  return {
    'train/mll': -loss.detach().item(),
  }
from copy import deepcopy

toy_model = SimplexGPModel(train_x, train_y).to(device)
# toy_model = KeOpsModel(train_x, train_y).to(device)
toy_mll = gp.mlls.ExactMarginalLogLikelihood(toy_model.likelihood, toy_model)
optimizer = torch.optim.Adam(toy_model.parameters(), lr=0.1)

for i in tqdm(range(100)):
    with torch.no_grad():
        toy_state_dict = deepcopy(toy_model.state_dict()) ## clone before it changes the reference.
    print(train(train_x, train_y, toy_model, toy_mll, optimizer))<jupyter_output><empty_output><jupyter_text>AutogradThis currently uses the approximation to the gradient, defined as another collection of filtering operations.<jupyter_code># with gp.settings.cg_tolerance(1e-2), \
#     gp.settings.max_preconditioner_size(100), \
#     gp.settings.max_root_decomposition_size(100):
#     model = BilateralGPModel(train_x, train_y, nu=1.5, order=1).to(device)
#     # model = KeOpsModel(train_x, train_y, nu=1.5).to(device)
#     model.base_covar_module.lengthscale = 1.0
#     mll = gp.mlls.ExactMarginalLogLikelihood(model.likelihood, model)

#     loss = -mll(model(train_x), train_y)
#     print(f'Loss: {loss.item()}')

#     loss.backward()

# for idx, p in enumerate(model.parameters()):
#     print(f'[{idx}] {p} ---> {p.grad}')<jupyter_output><empty_output><jupyter_text>Finite Difference<jupyter_code># import torch.nn.functional as F
# import torch.nn as nn

# def f(raw_ell):
#     with gp.settings.cg_tolerance(1e-2), \
#         gp.settings.max_preconditioner_size(100), \
#         gp.settings.max_root_decomposition_size(100), torch.no_grad():
#         model = BilateralGPModel(train_x, train_y, nu=1.5, order=1).to(device)
#         model.base_covar_module.raw_lengthscale = nn.Parameter(raw_ell)

#         mll = gp.mlls.ExactMarginalLogLikelihood(model.likelihood, model)
#         loss = -mll(model(train_x), train_y)

#     print(f'Raw Lengthscale: {model.base_covar_module.raw_lengthscale}; Loss: {loss}')
#     return loss

# raw_ell = model.base_covar_module.raw_lengthscale
# eps = 1e-4
# grads = []

# for _ in range(5):
#     grad = (f(raw_ell + eps) - f(raw_ell - eps)) / (2. * eps)
#     print(f'Finite Diff: {grad}')
#     grads.append(grad.item())
#     torch.cuda.empty_cache()

# print(f'{np.mean(grads)} +/- {2 * np.std(grads)}')<jupyter_output><empty_output><jupyter_text>JVP Checks<jupyter_code># from gpytorch.kernels.keops import RBFKernel, MaternKernel
# from gpytorch_lattice_kernel import MaternLattice, RBFLattice

# # K_gt = MaternKernel(nu=1.5).to(device)
# # f_gt = lambda x, y: K_gt(x, x) @ y

# K_lattice = RBFLattice(order=1).to(device)
# f_lattice = lambda x, y: (K_lattice(x, x) @ y).sum()
# train_x = torch.randn(5, 1).to(device).requires_grad_(True)
# train_y = torch.rand(5, 1).to(device).requires_grad_(False)
# torch.autograd.gradcheck(f_lattice, (train_x, train_y), eps=1e-4, rtol=1e-2, atol=1e-2)<jupyter_output><empty_output><jupyter_text>Toy Finite-Diff Gradients at Convergence<jupyter_code>import torch.nn as nn

def f(raw_ell, state_dict):
    with gp.settings.cg_tolerance(1e-2), \
        gp.settings.max_preconditioner_size(100), \
        gp.settings.max_root_decomposition_size(100), torch.no_grad():
        model = BilateralGPModel(train_x, train_y, nu=1.5, order=1).to(device)
        model.load_state_dict(state_dict)
        model.base_covar_module.raw_lengthscale = nn.Parameter(raw_ell)

        print(list(model.parameters()))

        mll = gp.mlls.ExactMarginalLogLikelihood(model.likelihood, model)
        loss = -mll(model(train_x), train_y)

    print(f'Raw Lengthscale: {model.base_covar_module.raw_lengthscale}; Loss: {loss}')
    return loss

with torch.no_grad():
    raw_ell = toy_state_dict['covar_module.base_kernel.raw_lengthscale']
    eps = 1e-5
    grads = []

    for _ in range(5):
        grad = (f(raw_ell + eps, toy_state_dict) - f(raw_ell - eps, toy_state_dict)) / (2. * eps)
        print(f'Finite Diff: {grad}')
        grads.append(grad.item())
        torch.cuda.empty_cache()

    print(f'{np.mean(grads)} +/- {2 * np.std(grads)}')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15142150803461063, 'fraction_numerical': 0.015039142974866089, 'mean_word_length': 4.282916213275299, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '573128', 'license': 'Apache-2.0', 'n_tokens_mistral': 1964, 'n_tokens_neox': 1889, 'n_words': 290, 'path': 'notebooks/viz_grad.ipynb', 'repo_name': 'activatedgeek/simplex-gp', 'size': '4854'}