Support non-numpy array backends by ColmTalbot · Pull Request #886 · bilby-dev/bilby

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.

mj-will

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

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

+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?

I agree. When we get close to merging I'll take this out.

Does this now need to be taken out? (Reading the mattermost it seems the priority is to get this merged)

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

        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.

Would this be a candidate for a small patch that uses the TF version for jax until jax supports it natively?

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

            )
        )
+
+    betaln,


Not sure what this is.

Not anything good.

Suggested change

betaln,

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

+        # 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

-            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

+    # 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

            )
        )
+
+    betaln,


Not anything good.

Suggested change

betaln,

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

+        # 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

-            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

+    # 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

        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.

Does this need to be actioned?

GregoryAshton

Okay, I got through about 60% of the diff and I'm pausing here so will submit the questions so far.

GregoryAshton · 2026-02-19T14:21:12Z

+from .utils import BackendNotImplementedError
+
+
+def erfinv_import(xp):


All of these functions would benefit from a docstring to explain they do the import given the type of array backend.

Done (for the one remaining function)

GregoryAshton · 2026-02-19T14:28:25Z

-            _cdf[val >= self.minimum] = 1. - np.exp(-val[val >= self.minimum] / self.mu)
-        return _cdf
+        with np.errstate(divide="ignore"):
+            return -val / self.mu - xp.log(xp.asarray(self.mu)) + xp.log(val >= self.minimum)


Ah okay - are the bounds being implemented here? But, I don't see the upper bound being implemented.

I think this is carried over from the existing implementation.

GregoryAshton · 2026-02-19T14:34:23Z


-            signal[mode] = waveform_polarizations[mode] * det_response
-        signal_ifo = sum(signal.values()) * mask
+            signal[mode] = waveform_polarizations[mode] * mask * det_response


It looks like this is changing the way the mask is being used. From operating on a view to operating on the full array but zeroing the False cases. Is that correct?

I think it isn't, and I've just moved this multiplication by the mask up by a line.
I'm not sure why, but I don't think it should make a big difference.

ColmTalbot · 2026-05-14T22:13:10Z

Python 3.10 doesn't have support for a vmappable version of logsumexp through scipy leading to this job failing (https://github.com/bilby-dev/bilby/actions/runs/25883935510/job/76070707573?pr=886).

How do people feel about dropping support for Python 3.10 in Bilby 3.10? Numpy dropped support about a year ago.

…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

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 force-pushed the bilback branch from 2d28818 to 230f623 Compare January 23, 2026 16:51

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 mentioned this pull request Feb 17, 2026

FEAT: Add precessing spin transformation #1044

Draft

GregoryAshton reviewed Feb 19, 2026

View reviewed changes

mj-will removed the to discuss To be discussed on an upcoming call label May 14, 2026

ColmTalbot and others added 10 commits May 14, 2026 22:18

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

1e3f4af

…likelihood

FEAT: support multiband and relative binning likelihoods

cf5c611

FEAT: make more conversions backend agnostic

2bfc833

FEAT: use more normal conversions

d29b860

FEAT: move backend switching code to bilby

c5eb323

FEAT: make core prior backend agnostic

9bec666

FEAT: make non-numpy arrays serializable

0d9aba6

BUG: fix some array conversion methods

aea0ea8

DEV: some more prior agnosticism

8e61b7f

TEST: make all prior tests run

b658a02

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

ColmTalbot added 25 commits May 15, 2026 13:35

FMT: fix formatting

fe76a67

BUG: fix bugs in testing

faf7535

Fix some more conversions

640d911

Add pytorch core testing

73f89b4

FMT: run precommits

bbb72d9

Make torch fully tested

fcfabdc

FMT: pre-commit fix

9b3b5b8

TEST: fix torch roq tests

12ba0b5

CI: prioritize torch tests

07a5ebe

TEST: another attempt to fix torch tests

42d8b07

Another attempt at fixing torch ROQ tests

c669bfa

Fix arrays of data setting

70f029d

BUG: fix some more roq array issues

0184120

Make ROQ calculations use correct array backend

b2cf8aa

BUG: fix a missing array case

86763f8

FMT: pre-commit fixes

a44cdbf

CI: drop torch tests for python 3.10

2542fa2

FMT: precommit fix

a68bd37

TEST: exclude studentt tests for jax

4e3ca65

Add some more explicit array casts

a909e24

BUG: bug fixes for prior and gw likelihoods

9343b45

BUG: fix array namespace for torch

b99fc35

Update patches and backend documentation

a7fbff1

Address some comments and add some docstrings

933dd94

Fix precommits

5ffa0d2

ColmTalbot force-pushed the bilback branch from b4ac884 to 5ffa0d2 Compare May 15, 2026 13:36

ColmTalbot added 4 commits May 15, 2026 13:40

MAINT: Don't track uv lock file

618ca4f

MAINT: drop python 3.10 support

607059b

Remove extra multibanding time marginalization lines

5254b5f

TEST: fix test failures

310b69a

                           )
                       )
+                  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)

		from .utils import BackendNotImplementedError


		def erfinv_import(xp):

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

GregoryAshton May 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

Uh oh!

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!

GregoryAshton left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ColmTalbot commented Jan 7, 2025 •

edited

Loading

GregoryAshton May 15, 2026 •

edited

Loading