LLNL / MuyGPyS

A fast, pure python implementation of the MuyGPs Gaussian process realization and training algorithm.
Other
25 stars 11 forks source link

Feature/heteroscedasticity #109

Closed alecmdunton closed 1 year ago

alecmdunton commented 1 year ago

Should have called my branch iss/70. This PR adds the heteroscedastic noise feature. Does not support the functions in MuyGPyS/examples. We should discuss the changes I made to MuyGPyS/gp/tensors: the function make_noise_tensor currently takes a NN_Wrapper object as input which (I think?) we may not want.

bwpriest commented 1 year ago

@alecmdunton tests/backend/mpi_correctness.py is breaking because you aren't chunking the distributed tensors correctly. You should run the test locally to debug rather than pushing changes to debug.

Similarly, tests/backend/torch_correctness.py is breaking because size is used incorrectly. Do you really need that function? Does shape not do what you need?

alecmdunton commented 1 year ago

I'm running into issues chunking the heteroscedastic noise tensor. Since we're treating it as a MuyGPs hyperparameter, are we actually able to chunk it while still using it to construct a model?

bwpriest commented 1 year ago

are we actually able to chunk it while still using it to construct a model?

Yes. The chunking happens within MuyGPyS.gp.tensors (check MuyGPyS._src.gp.tensors.pairwise_tensor, for example). All of the tensors whose first dimension is batch_count or test_count get chunked in the same way, such that the local tensors agree with each other, e.g., pairwise_diffs[i,:,:,:] refers to the same element as crosswise_diffs[i,:,:]. Done correctly, the same will be true of muygps.eps()[i,:,:].

This is why it is so important that we only use the index tensors (which index into [0, ..., batch_count - 1]) to construct tensors using MuyGPyS.gp.tensors, and subsequently stick to those tensors.

Let me know if you need more help understanding how this works.