Support non-numpy array backends #886

ColmTalbot · 2025-01-07T19:38:30Z

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:

prior sampling/rescale shapes - related to Conserve rescale shape #863
some priors won't return floats on float input

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 · 2026-01-23T15:48:31Z

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.

…likelihood

This required making some changes to the tests for conditional dicts as I've changed the output types and the backend introspection doesn't work on dict_items for some reason

mj-will

Some initial comments but I'll need to have another look.

mj-will · 2026-01-27T15:52:27Z

bilby/compat/patches.py

+    elif aac.is_cupy_namespace(xp):
+        from cupyx.scipy.special import erfinv
+    else:
+        raise BackendNotImplementedError


I think it would be useful to include the backend in the error.

mj-will · 2026-01-27T15:53:25Z

bilby/compat/utils.py

+__all__ = ["array_module", "promote_to_array"]
+
+
+def array_module(arr):


This would benefit from a doc-string

mj-will · 2026-01-27T15:54:21Z

bilby/compat/utils.py

+            return np
+
+
+def promote_to_array(args, backend, skip=None):


Have you thought about how devices would be handled here?

Moving arrays to a from GPUs can sometimes require more than just calling array.

mj-will · 2026-01-27T15:54:37Z

bilby/compat/utils.py

+            return np
+
+
+def promote_to_array(args, backend, skip=None):


Suggest adding a doc-string

mj-will · 2026-01-27T15:55:13Z

bilby/core/prior/analytical.py

+import os
+
 import numpy as np
+os.environ["SCIPY_ARRAY_API"] = "1"  # noqa  # flag for scipy backend switching


I worry slightly about having this hard coded. Does it introduce more overhead when using just numpy?

mj-will · 2026-01-27T16:01:03Z

bilby/core/prior/analytical.py

        This maps to the inverse CDF. This has been analytically solved for this case.
        """
-        return gammaincinv(self.k, val) * self.theta
+        return xp.asarray(gammaincinv(self.k, val)) * self.theta


Does this mean this is falling back to numpy?

Yeah, I should update/recheck this, but at least jax doesn't have good support for this, but it looks like tensorflow has a version that numpyro uses (jax-ml/jax#5350). cupy does have this function, so this workaround may have just been for jax. I could add a BackendNotImplementedError.

mj-will · 2026-01-27T16:04:22Z

bilby/core/prior/analytical.py

            )
        )
+
+    betaln,


Not sure what this is.

Not anything good.

Suggested change

betaln,

mj-will · 2026-01-27T16:06:09Z

bilby/core/prior/dict.py

+        # return self.check_ln_prob(sample, ln_prob,
+        #                           normalized=normalized)


Is the removal of this intentional?

I'm fairly sure it was, but I'll double check. I think check_ln_prob was problematic in some way.

mj-will · 2026-01-27T16:09:37Z

bilby/core/prior/dict.py

-            self[key].least_recently_sampled = result[key]
-            if isinstance(self[key], JointPrior) and self[key].dist.distname not in joint:
-                joint[self[key].dist.distname] = [key]
-            elif isinstance(self[key], JointPrior):
-                joint[self[key].dist.distname].append(key)
-        for names in joint.values():
-            # this is needed to unpack how joint prior rescaling works
-            # as an example of a joint prior over {a, b, c, d} we might
-            # get the following based on the order within the joint prior
-            # {a: [], b: [], c: [1, 2, 3, 4], d: []}
-            # -> [1, 2, 3, 4]
-            # -> {a: 1, b: 2, c: 3, d: 4}
-            values = list()
-            for key in names:
-                values = np.concatenate([values, result[key]])
-            for key, value in zip(names, values):
-                result[key] = value
-
-        def safe_flatten(value):
-            """
-            this is gross but can be removed whenever we switch to returning
-            arrays, flatten converts 0-d arrays to 1-d so has to be special
-            cased
-            """
-            if isinstance(value, (float, int)):
-                return value


Is removing this intentional?

Yeah, this is in line with one of the other open PRs to update this logic. I'll dig it out in my next pass.

mj-will · 2026-01-27T16:24:44Z

bilby/gw/utils.py

+    # delta_x = ifos[0].geometry.vertex - ifos[1].geometry.vertex
+    # theta, phi = zenith_azimuth_to_theta_phi(zenith, azimuth, delta_x)


Suggest we remove this.

Suggested change

# delta_x = ifos[0].geometry.vertex - ifos[1].geometry.vertex

# theta, phi = zenith_azimuth_to_theta_phi(zenith, azimuth, delta_x)

ColmTalbot

Thanks for the initial comments @mj-will I'll take a pass at them ASAP.

ColmTalbot · 2026-01-28T06:57:01Z

bilby/core/prior/analytical.py

            )
        )
+
+    betaln,


Not anything good.

Suggested change

betaln,

ColmTalbot · 2026-01-28T06:58:29Z

bilby/core/prior/analytical.py

        """
        at_peak = (val == self.peak)
-        return np.nan_to_num(np.multiply(at_peak, np.inf))
+        return at_peak * 1.0


ColmTalbot · 2026-01-28T07:00:18Z

bilby/core/prior/analytical.py

+            (xp.sin(val) - xp.sin(self.minimum)) /
+            (xp.sin(self.maximum) - xp.sin(self.minimum))
        )
+        _cdf *= val >= self.minimum


This kind of in-place operation works, it's just operations on slices that don't work. Things like patterns that sometimes existed

_cdf = ... _cdf[val < self.minimum] = 0 _cdf[val > self.maximum] = 1

ColmTalbot · 2026-01-28T07:01:44Z

bilby/core/prior/dict.py

+        # return self.check_ln_prob(sample, ln_prob,
+        #                           normalized=normalized)


I'm fairly sure it was, but I'll double check. I think check_ln_prob was problematic in some way.

ColmTalbot · 2026-01-28T07:02:42Z

bilby/core/prior/dict.py

-            self[key].least_recently_sampled = result[key]
-            if isinstance(self[key], JointPrior) and self[key].dist.distname not in joint:
-                joint[self[key].dist.distname] = [key]
-            elif isinstance(self[key], JointPrior):
-                joint[self[key].dist.distname].append(key)
-        for names in joint.values():
-            # this is needed to unpack how joint prior rescaling works
-            # as an example of a joint prior over {a, b, c, d} we might
-            # get the following based on the order within the joint prior
-            # {a: [], b: [], c: [1, 2, 3, 4], d: []}
-            # -> [1, 2, 3, 4]
-            # -> {a: 1, b: 2, c: 3, d: 4}
-            values = list()
-            for key in names:
-                values = np.concatenate([values, result[key]])
-            for key, value in zip(names, values):
-                result[key] = value
-
-        def safe_flatten(value):
-            """
-            this is gross but can be removed whenever we switch to returning
-            arrays, flatten converts 0-d arrays to 1-d so has to be special
-            cased
-            """
-            if isinstance(value, (float, int)):
-                return value


Yeah, this is in line with one of the other open PRs to update this logic. I'll dig it out in my next pass.

ColmTalbot · 2026-01-28T07:03:23Z

bilby/gw/utils.py

+    # delta_x = ifos[0].geometry.vertex - ifos[1].geometry.vertex
+    # theta, phi = zenith_azimuth_to_theta_phi(zenith, azimuth, delta_x)


Suggested change

# delta_x = ifos[0].geometry.vertex - ifos[1].geometry.vertex

# theta, phi = zenith_azimuth_to_theta_phi(zenith, azimuth, delta_x)

ColmTalbot · 2026-01-28T07:03:54Z

bilby/gw/utils.py

        The natural logarithm of the bessel function
    """
-    return np.log(i0e(value)) + np.abs(value)
+    xp = array_module(value)


Comment to self: use xp_wrap here.

ColmTalbot added the enhancement New feature or request label Jan 7, 2025

ColmTalbot marked this pull request as draft January 7, 2025 19:38

ColmTalbot force-pushed the bilback branch from b902545 to af6881d Compare October 2, 2025 16:06

ColmTalbot force-pushed the bilback branch from 95020e5 to 0eeeaa5 Compare December 11, 2025 15:45

ColmTalbot force-pushed the bilback branch 2 times, most recently from ea348fa to 771a8a9 Compare January 22, 2026 17:00

ColmTalbot marked this pull request as ready for review January 23, 2026 15:24

ColmTalbot changed the title ~~DRAFT: Support non-numpy array backends~~ Support non-numpy array backends Jan 23, 2026

ColmTalbot added >100 lines refactoring to discuss To be discussed on an upcoming call labels Jan 23, 2026

ColmTalbot and others added 18 commits January 23, 2026 11:51

FEAT: enable backend switching for base gravitational-wave transient …

5faa432

…likelihood

FEAT: support multiband and relative binning likelihoods

208f227

FEAT: make more conversions backend agnostic

c4d9bdf

FEAT: use more normal conversions

cd78b34

FEAT: move backend switching code to bilby

68abf3c

FEAT: make core prior backend agnostic

47041a1

FEAT: make non-numpy arrays serializable

e816198

BUG: fix some array conversion methods

7a785c2

DEV: some more prior agnosticism

7ebf340

TEST: make all prior tests run

b558ea6

This required making some changes to the tests for conditional dicts as I've changed the output types and the backend introspection doesn't work on dict_items for some reason

DEV: move some jax functionality to compat

af8d604

REFACTOR: use array backend for ln_i0

5b5fa6b

make distance marginalizatio backend transparent

c52be69

DEV: some more prior dict array refactoring

025e3d5

fix jax logic for distance marginalization

47baf2c

improve efficiency of setting up multibanding

b501d83

make high-dimensional gaussians jax compatible

ca7e4f8

make cubic spline calibration work with jax backend

a8a9b98

ColmTalbot added 15 commits January 23, 2026 11:51

BUG: fix some gnarly conversion corner cases

93cda56

BUG: fix multiband likelihood

0316acc

BUG: fix bug in array_namespace check

2c3f8fb

TEST: make sure healpix prior doesn't store state between calls

5d10a8a

FMT: example formatting fixes

27f2046

BUG: make sure indices don't overflow in roq

91ee508

BUG: fix multiband time marginalization setup

5ddf3e3

BUG: fix roq interpolation for out of bounds sample

79ae333

TYPO: fix typo in jax example

63b6f30

REFACTOR: refactor more roq likelihood tests

23a3d79

MAINT: revert new conversions

311ced4

CI: fix selecting only non-windows os

cb9703a

MAINT: make sure compat subpackages are listed in pyproject

ad23f4f

TYPO: Fix package list formatting in pyproject.toml

5930568

BUG: readd erroneously removed line

230f623

ColmTalbot force-pushed the bilback branch from 2d28818 to 230f623 Compare January 23, 2026 16:51

DOC: remove extraneous docstring

f65e668

mj-will added this to the 3.0.0 milestone Jan 27, 2026

mj-will reviewed Jan 27, 2026

View reviewed changes

ColmTalbot commented Jan 28, 2026

View reviewed changes

ColmTalbot added 10 commits January 28, 2026 11:50

Merge branch 'main' into bilback

6488bdb

TEST: fix test failures

2213038

TEST: start adding jax tests

a67b4ae

CI: add jax tests to CI

080df9d

Merge branch 'main' into bilback

2f9cd61

MAINT: add jax extras option

164bc70

Some more jax testing updates

2205fc2

MAINT: actually add jax requirements

aec63af

CI: don't trivially skip all tests...

cc79c54

Initial pass at making grid work with jax

9d4e01a

		__all__ = ["array_module", "promote_to_array"]


		def array_module(arr):

                           )
                       )
+                  betaln,

		# return self.check_ln_prob(sample, ln_prob,
		# normalized=normalized)

		# delta_x = ifos[0].geometry.vertex - ifos[1].geometry.vertex
		# theta, phi = zenith_azimuth_to_theta_phi(zenith, azimuth, delta_x)

Support non-numpy array backends #886

Are you sure you want to change the base?

Support non-numpy array backends #886

Uh oh!

Conversation

ColmTalbot commented Jan 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ColmTalbot commented Jan 23, 2026

Uh oh!

mj-will left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ColmTalbot left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ColmTalbot commented Jan 7, 2025 •

edited

Loading