MAINT: bump JAX to 0.7.2

[lucascolley]

some local CUDA failures:

================================================================================================ FAILURES =================================================================================================
_______________________________________________________________________________________ TestCov.test_empty[jax:gpu] _______________________________________________________________________________________

self = <tests.test_funcs.TestCov object at 0x742431e71c70>
xp = <module 'jax.numpy' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/jax/numpy/__init__.py'>

    @pytest.mark.xfail_xp_backend(Backend.JAX, reason="jax#32296")
    @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="sparse#877")
    def test_empty(self, xp: ModuleType):
        with warnings.catch_warnings(record=True):
            warnings.simplefilter("always", RuntimeWarning)
            warnings.simplefilter("always", UserWarning)
>           xp_assert_equal(
                cov(xp.asarray([], dtype=xp.float64)),
                xp.asarray(xp.nan, dtype=xp.float64),
            )

self       = <tests.test_funcs.TestCov object at 0x742431e71c70>
xp         = <module 'jax.numpy' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/jax/numpy/__init__.py'>

tests/test_funcs.py:530:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

actual = Array(-0., dtype=float64), desired = Array(nan, dtype=float64)

    def xp_assert_equal(
        actual: Array,
        desired: Array,
        *,
        err_msg: str = "",
        check_dtype: bool = True,
        check_shape: bool = True,
        check_scalar: bool = False,
    ) -> None:
        """
        Array-API compatible version of `np.testing.assert_array_equal`.

        Parameters
        ----------
        actual : Array
            The array produced by the tested function.
        desired : Array
            The expected array (typically hardcoded).
        err_msg : str, optional
            Error message to display on failure.
        check_dtype, check_shape : bool, default: True
            Whether to check agreement between actual and desired dtypes and shapes
        check_scalar : bool, default: False
            NumPy only: whether to check agreement between actual and desired types -
            0d array vs scalar.

        See Also
        --------
        xp_assert_close : Similar function for inexact equality checks.
        numpy.testing.assert_array_equal : Similar function for NumPy arrays.
        """
        xp = _check_ns_shape_dtype(actual, desired, check_dtype, check_shape, check_scalar)
        if not _is_materializable(actual):
            return
        actual_np = as_numpy_array(actual, xp=xp)
        desired_np = as_numpy_array(desired, xp=xp)
>       np.testing.assert_array_equal(actual_np, desired_np, err_msg=err_msg)
E       AssertionError:
E       Arrays are not equal
E
E       nan location mismatch:
E        ACTUAL: array(-0.)
E        DESIRED: array(nan)

actual     = Array(-0., dtype=float64)
actual_np  = array(-0.)
check_dtype = True
check_scalar = False
check_shape = True
desired    = Array(nan, dtype=float64)
desired_np = array(nan)
err_msg    = ''
xp         = <module 'jax.numpy' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/jax/numpy/__init__.py'>

src/array_api_extra/_lib/_testing.py:173: AssertionError
_________________________________________________________________________ TestIsClose.test_array_on_device_with_scalar[torch:gpu] _________________________________________________________________________

self = <tests.test_funcs.TestIsClose object at 0x742431ae53b0>
xp = <module 'array_api_compat.torch' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/array_api_compat/torch/__init__.py'>
device = device(type='cpu')

    def test_array_on_device_with_scalar(self, xp: ModuleType, device: Device):
        a = xp.asarray([0.01, 0.5, 0.8, 0.9, 1.00001], device=device)
        b = 1
        res = isclose(a, b)
        assert get_device(res) == device
>       xp_assert_equal(res, xp.asarray([False, False, False, False, True]))

a          = tensor([0.0100, 0.5000, 0.8000, 0.9000, 1.0000])
b          = 1
device     = device(type='cpu')
res        = tensor([False, False, False, False, False])
self       = <tests.test_funcs.TestIsClose object at 0x742431ae53b0>
xp         = <module 'array_api_compat.torch' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/array_api_compat/torch/__init__.py'>

tests/test_funcs.py:996:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

actual = tensor([False, False, False, False, False]), desired = tensor([False, False, False, False,  True], device='cuda:0')

    def xp_assert_equal(
        actual: Array,
        desired: Array,
        *,
        err_msg: str = "",
        check_dtype: bool = True,
        check_shape: bool = True,
        check_scalar: bool = False,
    ) -> None:
        """
        Array-API compatible version of `np.testing.assert_array_equal`.

        Parameters
        ----------
        actual : Array
            The array produced by the tested function.
        desired : Array
            The expected array (typically hardcoded).
        err_msg : str, optional
            Error message to display on failure.
        check_dtype, check_shape : bool, default: True
            Whether to check agreement between actual and desired dtypes and shapes
        check_scalar : bool, default: False
            NumPy only: whether to check agreement between actual and desired types -
            0d array vs scalar.

        See Also
        --------
        xp_assert_close : Similar function for inexact equality checks.
        numpy.testing.assert_array_equal : Similar function for NumPy arrays.
        """
        xp = _check_ns_shape_dtype(actual, desired, check_dtype, check_shape, check_scalar)
        if not _is_materializable(actual):
            return
        actual_np = as_numpy_array(actual, xp=xp)
        desired_np = as_numpy_array(desired, xp=xp)
>       np.testing.assert_array_equal(actual_np, desired_np, err_msg=err_msg)
E       AssertionError:
E       Arrays are not equal
E
E       Mismatched elements: 1 / 5 (20%)
E        ACTUAL: array([False, False, False, False, False])
E        DESIRED: array([False, False, False, False,  True])

actual     = tensor([False, False, False, False, False])
actual_np  = array([False, False, False, False, False])
check_dtype = True
check_scalar = False
check_shape = True
desired    = tensor([False, False, False, False,  True], device='cuda:0')
desired_np = array([False, False, False, False,  True])
err_msg    = ''
xp         = <module 'array_api_compat.torch' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/array_api_compat/torch/__init__.py'>

src/array_api_extra/_lib/_testing.py:173: AssertionError
____________________________________________________________________________________ TestPartition.test_nd[jax:gpu-2] _____________________________________________________________________________________

self = <tests.test_funcs.TestPartition object at 0x7424319298b0>
xp = <module 'jax.numpy' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/jax/numpy/__init__.py'>, ndim = 2

    @pytest.mark.parametrize("ndim", [2, 3, 4])
    def test_nd(self, xp: ModuleType, ndim: int):
>       self._test_nd(xp, ndim)

ndim       = 2
self       = <tests.test_funcs.TestPartition object at 0x7424319298b0>
xp         = <module 'jax.numpy' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/jax/numpy/__init__.py'>

tests/test_funcs.py:1531:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tests/test_funcs.py:1513: in _test_nd
    self._assert_valid_partition(x, k, y, xp, axis=i)
        base_shape = [2, 1]
        i          = 1
        k          = 1
        n          = 10
        ndim       = 2
        rng        = Generator(PCG64) at 0x741FC80F26C0
        self       = <tests.test_funcs.TestPartition object at 0x7424319298b0>
        shape      = [2, 10]
        x          = array([[5, 7, 5, 5, 2, 7, 2, 4, 6, 3],
       [9, 9, 7, 0, 8, 3, 0, 1, 7, 1]])
        xp         = <module 'jax.numpy' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/jax/numpy/__init__.py'>
        y          = Array([[ 2,  2,  7,  7,  6,  5,  5,  5,  4,  3],
       [ 0,  0,  8,  7, -8, -3,  1,  1, -7, -1]], dtype=int64)
        z          = array([[0.1485111 , 0.55676373],
       [0.68949362, 0.13056527]])
tests/test_funcs.py:1480: in _assert_valid_partition
    cls._assert_valid_partition(None, k, y[i, ...], xp, axis=axis)
        axis       = 0
        cls        = <class 'tests.test_funcs.TestPartition'>
        i          = 1
        k          = 1
        x_np       = array([[5, 7, 5, 5, 2, 7, 2, 4, 6, 3],
       [9, 9, 7, 0, 8, 3, 0, 1, 7, 1]])
        xp         = <module 'jax.numpy' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/jax/numpy/__init__.py'>
        y          = Array([[ 2,  2,  7,  7,  6,  5,  5,  5,  4,  3],
       [ 0,  0,  8,  7, -8, -3,  1,  1, -7, -1]], dtype=int64)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

cls = <class 'tests.test_funcs.TestPartition'>, x_np = None, k = 1, y = Array([ 0,  0,  8,  7, -8, -3,  1,  1, -7, -1], dtype=int64)
xp = <module 'jax.numpy' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/jax/numpy/__init__.py'>, axis = 0

    @classmethod
    def _assert_valid_partition(
        cls,
        x_np: np.ndarray | None,
        k: int,
        y: Array,
        xp: ModuleType,
        axis: int | None = -1,
    ):
        """
        x_np : input array
        k : int
        y : output array returned by the partition function to test
        """
        if x_np is not None:
            assert y.shape == np.partition(x_np, k, axis=axis).shape
        if y.ndim != 1 and axis == 0:
            assert isinstance(y.shape[1], int)
            for i in range(y.shape[1]):
                cls._assert_valid_partition(None, k, y[:, i, ...], xp, axis=0)
        elif y.ndim != 1:
            assert axis is not None
            axis = axis - 1 if axis != -1 else -1
            assert isinstance(y.shape[0], int)
            for i in range(y.shape[0]):
                cls._assert_valid_partition(None, k, y[i, ...], xp, axis=axis)
        else:
            if k > 0:
                assert xp.max(y[:k]) <= y[k]
>           assert y[k] <= xp.min(y[k:])
E           AssertionError: assert Array(0, dtype=int64) <= Array(-8, dtype=int64)
E            +  where Array(-8, dtype=int64) = <function min at 0x7422c40e4b80>(Array([ 0,  8,  7, -8, -3,  1,  1, -7, -1], dtype=int64))
E            +    where <function min at 0x7422c40e4b80> = <module 'jax.numpy' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/jax/numpy/__init__.py'>.min

axis       = 0
cls        = <class 'tests.test_funcs.TestPartition'>
k          = 1
x_np       = None
xp         = <module 'jax.numpy' from '/home/lucas/ghq/github.com/data-apis/array-api-extra/.pixi/envs/tests-cuda/lib/python3.13/site-packages/jax/numpy/__init__.py'>
y          = Array([ 0,  0,  8,  7, -8, -3,  1,  1, -7, -1], dtype=int64)

tests/test_funcs.py:1484: AssertionError

[lucascolley]

the Partition failures seem non-deterministic, won't bother investigating further right now