Skip to content

Conversation

@ColmTalbot
Copy link
Collaborator

@ColmTalbot ColmTalbot commented Jan 7, 2025

I've been working on this PR on and off for a few months, it isn't ready yet, but I wanted to share it in case other people had early opinions.

The goal is to make it easier to interface with models/samplers implemented in e.g., JAX, that support GPU/TPU acceleration and JIT compilation.

The general guiding principles are:

  • when possible maintain existing behaviour with numpy/builtin arguments
  • work introspectively so users don't need to specify the target backend, but use input types
  • write as little backend specific code as possible, mostly through using the array-api specification and scipy interoperability

The primary changes so far are:

  • making most priors backend independent, there are a few holdouts where the underlying scipy functionality isn't compatible yet
  • core likelihoods mostly work with data from any backend
  • GW likelihoods work with any backend supported by the source function
  • the GW detector objects don't work via introspection, they need to be manually set
  • GW geometry (currently in bilby_cython) is handled via multiple-dispatch and added back into bilby

Changed behaviour:

Remaining issues:

  • Saving/loading nun-numpy arrays in result files may not work
  • I added some additional parameter conversions that I will remove
  • the bilby.gw.jaxstuff file should be removed and relevant functionality be moved elsewhere, it's currently just used for testing
  • the ROQ likelihood hasn't been ported
  • add more testing with JAX
  • translate some of the hyperparameter functionality, c.f., GWPopulation

@ColmTalbot ColmTalbot added the enhancement New feature or request label Jan 7, 2025
@ColmTalbot ColmTalbot marked this pull request as draft January 7, 2025 19:38
@ColmTalbot ColmTalbot force-pushed the bilback branch 2 times, most recently from ea348fa to 771a8a9 Compare January 22, 2026 17:00
@ColmTalbot ColmTalbot marked this pull request as ready for review January 23, 2026 15:24
@ColmTalbot ColmTalbot changed the title DRAFT: Support non-numpy array backends Support non-numpy array backends Jan 23, 2026
@ColmTalbot ColmTalbot added >100 lines refactoring to discuss To be discussed on an upcoming call labels Jan 23, 2026
@ColmTalbot
Copy link
Collaborator Author

This is now ready for review.
There are some things that won't work with JAX at the moment, e.g., various combinations of likelihood marginalization/acceleration.
I think we should accept this at the moment, for at least a bilby v3 alpha/beta release, and keep chipping away at the various subcases over time.

There are a lot of changes, but most of them are essentially np -> xp.
Some things required refactoring to avoid modifying slices of arrays as JAX doesn't like that.

Bilby can once again be installed without bilby.cython.
This should improve our general portability, but when bilby_cython is installed it will be used.

I've managed to keep test changes minimal:

  • I updated the joint prior test to make it more stringent (keys more randomly ordered).
  • I refactored some expensive prior initialization that was dramatically slowing things down.
  • I improved the logic for figuring out when ROQs are available to help my local testing.
  • Some mocks of numpy had to be updated.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request refactoring to discuss To be discussed on an upcoming call >100 lines

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant