Allow testing methods to compare arrays of different backends

I'm setting up a test where I would like to do

import array_api_extra as xpx
import jax.numpy as jnp
import numpy as np

xpx.testing.assert_equal(np.asarray([0]), jnp.asarray([0]), check_backends=False)

It would be nice to be able to explicitly check these examples, but should of course default to check_backends=True.

could you link to the use-case? We used to have this in the versions that were in SciPy, but when taking a second look while switching to xpx.testing we decided it was cleaner to just always avoid this pattern.

I'd be happy to add this if the use-case is compelling.

So for our specific use case we are wanting to compare the outputs of an inverse spherical harmonic transform. Our current code has a wrapper around https://pypi.org/project/healpy and https://pypi.org/project/healpix, so that they can handle any array backend by internally converting to NumPy.

We would like to use https://pypi.org/project/s2fft to compute out spherical harmonic transforms natively in JAX, so that we never leave the GPU. So we can't do this in the classic Array API way, and would have to do an if xp.__name__ == "jax.numpy" pattern.

Regardless, it would be nice to use array-api-extra.testing.assert_equal to compare the output of the NumPy and JAX transform. I guess we could just use numpy.testing.assert_array_equal for that bit.

Do you have xp available at assertion time? The thinking is that check_backends=False is no more concise than wrapping the desired array with xp.asarray.

Yeah this is what I'm currently doing, just looks a bit ugly!

def test_inverse_transform() -> None:
    jax = pytest.importorskip("jax", reason="test requires jax")

    alm_numpy = np.asarray([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], dtype=np.float32)
    alm_jax = jax.numpy.asarray(alm_numpy, dtype=jax.numpy.float32)
    lmax = 2
    nside = 2

    # https://github.com/data-apis/array-api-extra/issues/1005
    xpx.testing.assert_equal(
        glass.harmonics.inverse_transform(alm_numpy, lmax=lmax, nside=nside),
        np.asarray(glass.harmonics.inverse_transform(alm_jax, lmax=lmax, nside=nside)),
    )

what I meant was having the test parametrised by xp, then you have alm_actual as the first param, and xp.asarray(alm_desired) as the second param.

having the test parametrised by xp, then you have alm_actual as the first param, and xp.asarray(alm_desired) as the second param.

If I'm not wrong, this was the check_namespace mechanism in SciPy that we removed in the process of adapting xpx.testing, right?

yep

It might then just make sense to add it in extra. SciPy used to have it, and GLASS also seems to have a use case.

yeah but see my messages above, I'm still not convinced that it is cleaner than calling xp.asarray

what I meant was having the test parametrised by xp, then you have alm_actual as the first param, and xp.asarray(alm_desired) as the second param.

I'm not sure I follow. How would that test specifically a NumPy mechanism Vs JAX? glass.harmonics.inverse_transform does something different depending on input

You always compute the reference value with NumPy arrays and the actual value with xp arrays.

Then, for the assertion, the reference value is coerced with xp.asarray, being explicit about the expected namespace without changing the reference values.

Does that make sense?

I think so, so something like this?

@pytest.mark.skipif(not (HAVE_JAX and HAVE_S2FFT), reason="test requires jax and s2fft")
def test_inverse_transform(
    healpix_inputs: type[HealpixInputs],
    rng: UnifiedGenerator,
) -> None:
    import jax
    import jax.numpy as jnp

    alm = healpix_inputs.alm(rng=rng)

    with jax.enable_x64(True):  # noqa: FBT003
        expected = glass.harmonics.inverse_transform(
            alm,
            lmax=healpix_inputs.lmax,
            nside=healpix_inputs.nside,
        )
        actual = glass.harmonics.inverse_transform(
            jnp.asarray(alm),
            lmax=healpix_inputs.lmax,
            nside=healpix_inputs.nside,
        )
        xpx.testing.assert_close(actual, jnp.asarray(expected), atol=1e-12, rtol=0)

yeah, that's the idea. Do you think that would be significantly cleaner with check_backends=False?

Well not significantly of course. Just feels a bit similar to me.

I'd err on keeping this as-is for now then, but do let me know if a more compelling use-case appears

Navigation Menu

Repository navigation

Allow testing methods to compare arrays of different backends #1005

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Allow testing methods to compare arrays of different backends #1005

Description

Activity

lucascolley commented on Sep 30, 2026

paddyroddy commented on Sep 30, 2026

lucascolley commented on Sep 30, 2026

paddyroddy commented on Oct 1, 2026

lucascolley commented on Oct 1, 2026

prady0t commented on Oct 1, 2026

lucascolley commented on Oct 1, 2026

prady0t commented on Oct 1, 2026

lucascolley commented on Oct 1, 2026

paddyroddy commented on Oct 1, 2026

lucascolley commented on Oct 1, 2026

paddyroddy commented on Oct 2, 2026

lucascolley commented on Oct 2, 2026

paddyroddy commented on Oct 2, 2026

lucascolley commented on Oct 2, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions