Skip to content

Gaussian Processes via JAX - #129

Merged
asher-m merged 22 commits into
masterfrom
dev/gp-jax
Oct 1, 2026
Merged

asher-m merged 22 commits into
masterfrom
dev/gp-jax

Conversation

@asher-m

@asher-m asher-m commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

This PR primarily adds Gaussian processes via JAX.

We'll discuss added features and features remaining to be added in our meeting.

* pending commit on MWE
  * tutorial works (all features, all dimensions), but it's kinda messy
  * want to clean up a few things, then happy to commit
* small changes to gp and _equinox modules
  * bug fixes
  * some changes to facilitate real use-cases
* still want to implement dynamic selection of next sampling coordinate
* still want to change some symbol names, but works
* more coming for docs
* more coming for different noise structures, too
…namespace

* may break this out further later
* for now, I don't want any _CF* objects getting confused with covariance kernels
* will want to convert formalism/gp.md to .rst eventually
* (cross)linking in .md is more terse than .rst
@asher-m
asher-m requested a review from akilnarayan September 10, 2026 21:35
@asher-m asher-m self-assigned this Sep 10, 2026
* need to reorganize tutorial docs
* reverted python 3.9 breaking changes on rebase
* need to find solution for copying demos/sources from their permanent home into docs at build
@asher-m

asher-m commented Sep 21, 2026

Copy link
Copy Markdown
Collaborator Author

I reverted the introduction of a Sphinx addon requiring Python version >= 3.10, but that exposed another pathology: Python 3.9, 3.10, and 3.11 are deprecated by JAX and no longer receive updates, or at least they are not shown as supported by the newest version of JAX on PyPI.

I figure we have a few reasonable choices:

  1. Pin JAX to a sufficiently old version such that our use of the JAX API is constant across the different versions pulled after dependency resolution across our supported versions of Python (i.e., 3.9, 10, and 11), e.g. pin JAX to 0.4(.x) assuming that supports our Python.
    • This could work, but I'll need to do some tinkering to ensure the new functionality works.
  2. Deprecate 3.9, 10, and 11 (!?). This seems to be the direction of the broader Python ecosystem. We've briefly discussed bumping our supported Python up to something newer (3.14, 15?) but haven't yet done so.
    • Feels silly and/or undesirable to deprecate all versions of Python we currently support, but this feels like the direction things are going, anyway.
  3. Something else?

Not sure the direction we want to go. Both of these seem reasonable with their respective tradeoffs.

@akilnarayan, let me know what you think, or we can discuss this tomorrow.

@akilnarayan akilnarayan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed w/ asher-m

@asher-m
asher-m merged commit 880cd17 into master Oct 1, 2026
6 of 14 checks passed
@asher-m
asher-m deleted the dev/gp-jax branch October 1, 2026 21:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants