Kernels not compatible with torch.func.functional_call?

[bchen0]

I want to use torch.func.functional_call to call kernels with different parameters.

The "normal" way would be to change the state dictionary directly.

torch.manual_seed(0)
kernel = RBFKernel()
test_array = torch.randn(2,2)
print(kernel.forward(test_array, test_array))

state_dict = copy.deepcopy(kernel.state_dict())
state_dict['raw_lengthscale'] = torch.Tensor([[1.]])
kernel.load_state_dict(state_dict)
print(kernel.forward(test_array, test_array))

This prints a matrix with 2.57e-7 in the off diagonal in the first case and 0.0146 in the second. However, if I try to do this with torch.func.functional_call as follows:

state_dict_copy = copy.deepcopy(kernel.state_dict())
def functional_call(x):
    state_dict_copy['raw_lengthscale'] = x
    return torch.func.functional_call(
        kernel,
        state_dict_copy,
        args=(test_array, test_array)
    )
print(functional_call(torch.Tensor([[1.]])).to_dense())
print(functional_call(torch.Tensor([[0.]])).to_dense())

this prints a matrix with 2.57e-7 in the off diagonal in both cases.

Any suggestions on why this does not work? One thing that is interesting to note is that the output of functional_call is a gpytorch.lazy.lazy_evaluated_kernel_tensor.LazyEvaluatedKernelTensor object instead of just a tensor - perhaps that is related?

The reason I want to use the second method instead of the first is that I have a mapping x --> lengthscales and I want to take the derivative of the output of a GP wrt x.

Thanks!

Edit: I have also tried this with _reparametrize_module as suggested here: #2662 (comment), but this also yields the same (incorrect) results

from torch.nn.utils.stateless import _reparametrize_module

torch.manual_seed(0)
kernel = RBFKernel()
test_array = torch.randn(2,2)

state_dict_copy = copy.deepcopy(kernel.state_dict())
def functional_call(x):
    state_dict_copy['raw_lengthscale'] = x
    with _reparametrize_module(kernel, state_dict_copy):
        result = kernel(test_array, test_array)
    return result

print(functional_call(torch.Tensor([[17.]])).to_dense())
print(functional_call(torch.Tensor([[0.]])).to_dense())

[gpleiss]

Because GPyTorch uses the linear_operator library under-the-hood, it's not well designed for torch's functional library.
You could try running all of this under the context manager with gpytorch.settings.lazily_evaluate_kernels(False), though I'm not 100% sure it would fix the issue.