Feature Request

Given a collection of GPJax Dataset objects, it would be nice to have functionality to create a HeterotopicDataset object.

Hello! I am currently looking at this repo, and finding some issues in some of the files. As I think I would like to utilize a lot of this functionality, I am "fixing" things as I find them.

I currently have a PR ready to go that fixes the gplvm.pct.py example, but cannot push it for code review, as I do not have access rights.

perhaps you can grant me a temporary token so I can push this (and other) PRs?

Bug Report

MOGPJax version:
0.0.2
Current behavior:

when running pytest on clean virtual environment (using python 3.10.0 and running pip install -e .), every test fails (144 of them) (and there are 792 warnings).

Expected behavior:

All tests pass.

Steps to reproduce:

See above: clean virtualenv running python 3.10.0 and running pip install -e . and pytest.

Now, if I comment out the line:

assert isinstance(mll(params), jnp.DeviceArray)

then all of the tests pass.

Inspecting the return type of mll(params) with type(mll(params)) returns jaxlib.xla_extension.Array, but we cannot
replace the test with

assert isinstance(mll(params), jaxlib.xla_extension.Array)

since

AttributeError: module 'jaxlib.xla_extension' has no attribute 'Array'

Other information:

I can suppress the warnings by replacing (in test_gps.py):

@pytest.mark.parametrize("kernel", [gpx.Matern12, gpx.Matern32, gpx.Matern52, gpx.RBF])

with

@pytest.mark.parametrize("kernel", [jk.Matern12, jk.Matern32, jk.Matern52, jk.RBF])

wherein I import jaxkern.kernels as jk