ENH: _lib._util: make _lazywhere compatible with Array API

[mdhaber]

Reference issue

Along the lines of gh-18668

What does this implement/fix?

This makes scipy._lib._util compatible with the array API. Also:

• The output dtype is now determined based on the types required by the outputs (rather than the types of the inputs).
• The input validation s is simplified.
• ~cond is now computed only once.
• np.empty is now used to initialize the output array when f2 is specified (instead of np.full).

For small (NumPy) arrays (<10k elements), this version tends to be faster. For moderate arrays (say 1 million elements), the execution time depends a lot on cond. Its distribution is spread out over more than order of magnitude, but is slightly (<6%) slower on average (gmean) with this PR. I'd be happy to find out that this can be improved, but I've studied it, and it seems to be accounted for by the use of array API-compatible features. Fortunately, it seems to asymptote to <7% slower as the array size increases.

Temporary memory is increased, but that seems necessary if we want to determine the output type based on what is going to be output (rather than the type of the input).

Additional information

Todo:

• unit tests (especially with array API array)
• When inputs are scalars, _lazywhere should return a scalar rather than 0D array. Is there a canonical way to determine the difference between the two (since the size, shape, and ndim are the same)? np.isscalar, of course
• The output is an ndarray even when the input is an array_api.Array. That's because xp.full and xp.empty create ndarrays. I'm assuming that's supposed to happen for now?

[mdhaber]

Once gh-18927 merges, I'll add this test:

class TestLazywhere:
    arrays = npst.arrays(dtype=npst.floating_dtypes(), shape=npst.array_shapes())
    fillvalue = npst.arrays(dtype=npst.floating_dtypes(), shape=tuple())
    p = strategies.floats(min_value=0, max_value=1)

    @array_api_compatible
    @given(arrays=arrays, fillvalue=fillvalue, p=p)
    def test_basic(self, arrays, fillvalue, p, xp):
        arrays = xp.asarray(arrays)
        fillvalue = xp.asarray(fillvalue)

        def f(*args):
            return args[0]

        def f2(*args):
            return args[0]*0.5

        if arrays.ndim <= 1:
            arrays = (arrays,)

        rng = np.random.default_rng(84268954369357456)
        cond = xp.asarray(rng.random(size=arrays[0].shape) > p)

        res = _lazywhere(cond, arrays, f, fillvalue)
        ref = xp.where(cond, f(*arrays), fillvalue)
        assert_equal(res, ref)

        res = _lazywhere(cond, arrays, f, f2=f2)
        ref = xp.where(cond, f(*arrays), f2(*arrays))
        assert_equal(res, ref)

@tupui is this how you'd use @array_api_compatible here?

Note that https://scipy.github.io/devdocs/dev/api-dev/array_api.html doesn't say where to import that decorator from.
This PR will add

  from scipy.conftest import array_api_compatible
  ...

to the top of the example on that page.

I'm not going to add a test of the input validation. It was there, so I preserved it, but I would rather remove it than test it.

[tupui]

@tupui is this how you'd use @array_api_compatible here?

Yes 👍

Also maybe np.where should also be converted to xp.where.

doesn't say where to import that decorator from.
This PR will add

from scipy.conftest import array_api_compatible
...
to the top of the example on that page.

Good catch 👍 Yes this is the correct import.

[tylerjereddy]

I just played around with this branch + copy/paste of the test mentioned in case the feedback is helpful.

Testing with torch: python dev.py test -b pytorch -t scipy/stats/tests/test_distributions.py::TestLazywhere produces error below the fold:

_________________________________________________________________________________________________________ TestLazywhere.test_basic[torch] _________________________________________________________________________________________________________
scipy/stats/tests/test_distributions.py:9455: in test_basic
    @given(arrays=arrays, fillvalue=fillvalue, p=p)
        f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7fe19fb620c0>
        self       = <scipy.stats.tests.test_distributions.TestLazywhere object at 0x7fe19fa37610>
        xp         = <module 'torch' from '/home/treddy/python_venvs/py_311_scipy_dev/lib/python3.11/site-packages/torch/__init__.py'>
scipy/stats/tests/test_distributions.py:9472: in test_basic
    res = _lazywhere(cond, arrays, f, fillvalue)
        arrays     = (tensor([0.], dtype=torch.float16),)
        cond       = array([ True])
        f          = <function TestLazywhere.test_basic.<locals>.f at 0x7fe19fbcb4c0>
        f2         = <function TestLazywhere.test_basic.<locals>.f2 at 0x7fe19fbc9d00>
        fillvalue  = tensor(0., dtype=torch.float16)
        p          = 0.0
        rng        = Generator(PCG64) at 0x7FE19FC06EA0
        self       = <scipy.stats.tests.test_distributions.TestLazywhere object at 0x7fe19fa37610>
        xp         = <module 'torch' from '/home/treddy/python_venvs/py_311_scipy_dev/lib/python3.11/site-packages/torch/__init__.py'>
scipy/_lib/_util.py:87: in _lazywhere
    xp = array_namespace(cond, *arrays)
        arrays     = (tensor([0.], dtype=torch.float16),)
        cond       = array([ True])
        f          = <function TestLazywhere.test_basic.<locals>.f at 0x7fe19fbcb4c0>
        f2         = None
        fillvalue  = tensor(0., dtype=torch.float16)
scipy/_lib/_array_api.py:95: in array_namespace
    return array_api_compat.array_namespace(*arrays)
        arrays     = [array([ True]), tensor([0.], dtype=torch.float16)]
scipy/_lib/array_api_compat/array_api_compat/common/_helpers.py:110: in array_namespace
    raise TypeError(f"Multiple namespaces for array inputs: {namespaces}")
E   TypeError: Multiple namespaces for array inputs: {<module 'scipy._lib.array_api_compat.array_api_compat.numpy' from '/home/treddy/github_projects/scipy/build-install/lib/python3.11/site-packages/scipy/_lib/array_api_compat/array_api_compat/numpy/__init__.py'>, <module 'scipy._lib.array_api_compat.array_api_compat.torch' from '/home/treddy/github_projects/scipy/build-install/lib/python3.11/site-packages/scipy/_lib/array_api_compat/array_api_compat/torch/__init__.py'>}
E   Falsifying example: test_basic(
E       self=<scipy.stats.tests.test_distributions.TestLazywhere object at 0x7fe19fa37610>,
E       xp=<module 'torch' from '/home/treddy/python_venvs/py_311_scipy_dev/lib/python3.11/site-packages/torch/__init__.py'>,
E       arrays=array([0.], dtype=float16),
E       fillvalue=array(0., dtype=float16),
E       p=0.0,
E   )

I get an almost identical error from testing with CuPy (python dev.py test -b cupy -t scipy/stats/tests/test_distributions.py::TestLazywhere).

Looks like, in your draft test, cond from np.random.. produces a NumPy array, so adding cond = xp.asarray(cond) seems to help a bit in my hands. I do note that mixing hypothesis in here does make things a bit tricky on top of array API. Testing still fails on the device (with CuPy) with that patch because we can't just use assert_equal(res, ref) with host-device transfers. I haven't checked if we've merged an official way to do that just yet, I know I did draft some ideas.

[mdhaber]

Thanks. Yes, I planned to test this once the Hypothesis PR merges, but it looks like adding xp.asarray to the cond variable will save me some trouble.

I do note that mixing hypothesis in here does make things a bit tricky on top of array API. Testing still fails on the device (with CuPy) with that patch because we can't just use assert_equal(res, ref) with host-device transfers.

Is the problem inherent to Hypothesis, or is it that np.testing.assert_equal doesn't work for arbitrary array types?

[tupui]

Testing still fails on the device (with CuPy) with that patch because we can't just use assert_equal(res, ref) with host-device transfers. I haven't checked if we've merged an official way to do that just yet, I know I did draft some ideas.

My PR incorporated your code Tyler:

if "cupy" in xp.__name__:
    import cupy as cp
    cp.testing.assert_allclose(...)
else:
    np.testing.assert_allclose(...)

For the fft PR, I actually asked @lucascolley to make a standalone helper function that we can use elsewhere more easily.

It's a bit unfortunate that there is no testing bits in the standard itself. I feel like this would be helpful. Maybe we can push this at least in the compat library.

[tylerjereddy]

Edit: Pamphile clarified we already have some device testing shims

For torch I also noticed this, which you'll see when you start iterating I suppose:

scipy/_lib/_util.py:93: in _lazywhere
    cond, arrays = args[0].astype(bool, copy=False), args[1:]
E   AttributeError: 'Tensor' object has no attribute 'astype'

I was able to patch around that with:

--- a/scipy/_lib/_util.py
+++ b/scipy/_lib/_util.py
@@ -90,7 +90,7 @@ def _lazywhere(cond, arrays, f, fillvalue=None, f2=None):
         raise ValueError("Exactly one of `fillvalue` or `f2` must be given.")
 
     args = xp.broadcast_arrays(cond, *arrays)
-    cond, arrays = args[0].astype(bool, copy=False), args[1:]
+    cond, arrays = xp.astype(args[0], xp.bool, copy=False), args[1:]
 
     temp1 = f(*(arr[cond] for arr in arrays))

But then hypothesis started doing stuff that seemed hard for the torch control flow to handle:

scipy/stats/tests/test_distributions.py:9455: in test_basic
    @given(arrays=arrays, fillvalue=fillvalue, p=p)
        f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f1018bee0c0>
        self       = <scipy.stats.tests.test_distributions.TestLazywhere object at 0x7f1018ac7410>
        xp         = <module 'torch' from '/home/treddy/python_venvs/py_311_scipy_dev/lib/python3.11/site-packages/torch/__init__.py'>
scipy/stats/tests/test_distributions.py:9458: in test_basic
    fillvalue = xp.asarray(fillvalue)
E   ValueError: given numpy array has byte order different from the native byte order. Conversion between byte orders is currently not supported.
E   Falsifying example: test_basic(
E       self=<scipy.stats.tests.test_distributions.TestLazywhere object at 0x7f1018ac7410>,
E       xp=<module 'torch' from '/home/treddy/python_venvs/py_311_scipy_dev/lib/python3.11/site-packages/torch/__init__.py'>,
E       arrays=array([0.], dtype=float16),
E       fillvalue=array(0., dtype='>f2'),
E       p=0.0,
E   )

[mdhaber]

I see astype is a function in the Array API, not a method. Got it.

[tylerjereddy]

astype is in the array API. Shouldn't array_api_compat be taking care of that @tupui?

I think it is with the small patch I added

[tupui]

I see astype is a function in the Array API, not a method. Got it.

Yes that's the main design change, everything is a function instead of a method in the Array API standard.

[lucascolley]

For the fft PR, I actually asked @lucascolley to make a standalone helper function that we can use elsewhere more easily.

For my PR (which will appear on the main repo soon) this has been added to conftest.py:

def set_assert_allclose(xp=None):
    if xp is None:
        return npt.assert_allclose
    if 'cupy' in xp.__name__:
        import cupy as cp
        return cp.testing.assert_allclose
    return npt.assert_allclose

[lucascolley]

That's because xp.full and xp.empty create ndarrays.

Is this true? With xp == numpy.array_api they seem to create numpy.array_api arrays for me:

In [1]: import numpy as np

In [2]: import numpy.array_api as npa
<ipython-input-2-47d1a2ee1bbf>:1: UserWarning: The numpy.array_api submodule is still experimental. See NEP 47.
  import numpy.array_api as npa

In [3]: a = npa.asarray([1, 2, 3, 4])

In [4]: cond = a > 2

In [5]: isinstance(cond, npa._array_object.Array)
Out[5]: True

In [6]: isinstance(cond, np.ndarray)
Out[6]: False

In [7]: x = npa.full(cond.shape, fill_value=npa.nan)

In [8]: isinstance(x, npa._array_object.Array)
Out[8]: True

In [9]: isinstance(x, np.ndarray)
Out[9]: False

Perhaps xp is not being set correctly?

[mdhaber]

I don't think xp is set to numpy.array_api. IIUC, For now, we are using array_api_compat.numpy on purpose.

[lucascolley]

Very sorry if I am misunderstanding, but I think that xp should be set to numpy.array_api as long as the global switch SCIPY_ARRAY_API is not False.

The first line of _lazywhere is xp = array_namespace(cond, *arrays), which calls array_namespace in _array_api.py:

if not _GLOBAL_CONFIG["SCIPY_ARRAY_API"]:
        # here we could wrap the namespace if needed
        return array_api_compat_numpy

arrays = [array for array in arrays if array is not None]

compliance_scipy(*arrays) 

return array_api_compat.array_namespace(*arrays)

Which in turn calls array_namespace from array_api_compat:

namespaces = set()
    for x in xs:
        if isinstance(x, (tuple, list)):
            namespaces.add(array_namespace(*x, _use_compat=_use_compat))
        elif hasattr(x, '__array_namespace__'):
            namespaces.add(x.__array_namespace__(api_version=api_version))
        elif _is_numpy_array(x):
            _check_api_version(api_version)
            if _use_compat:
                from .. import numpy as numpy_namespace
                namespaces.add(numpy_namespace)
            else:
                import numpy as np
                namespaces.add(np)
        # ...
    if not namespaces:
        raise TypeError("Unrecognized array input")

    if len(namespaces) != 1:
        raise TypeError(f"Multiple namespaces for array inputs: {namespaces}")

    xp, = namespaces

    return xp

Since hasattr(x, '__array_namespace__') is True when x is a numpy.array_api._array_object.Array, numpy.array_api is returned by this function.

Am I missing something? It seems to me that xp would only be array_api_compat.numpy for a numpy.array_api._array_object.Array input if SCIPY_ARRAY_API was set to False, which shouldn't be the case whenever -b is used with python dev.py test.

[tupui]

Very sorry if I am misunderstanding, but I think that xp should be set to numpy.array_api as long as the global switch SCIPY_ARRAY_API is not False.

array_api_compat.numpy is returned when the flag is False. This is so that normal NumPy arrays are made compatible with the standard. Hence we can start adding compatibility.

When the flag is on. The namespace that we get back is coming from the the function array_api_compat.array_namespace. Here it's making the decision for NumPy to give us back the portable version, i.e. array_api_compat.numpy as well.

See here:

scipy/scipy/_lib/_array_api.py

Lines 87 to 95 in 73c8d7b

	if not _GLOBAL_CONFIG["SCIPY_ARRAY_API"]:
	# here we could wrap the namespace if needed
	return array_api_compat_numpy
	arrays = [array for array in arrays if array is not None]
	compliance_scipy(*arrays)
	return array_api_compat.array_namespace(*arrays)

[lucascolley]

Here it's making the decision for NumPy to give us back the portable version, i.e. array_api_compat.numpy as well.

I don't agree! As stated in my last comment:

Since hasattr(x, '__array_namespace__') is True when x is a numpy.array_api._array_object.Array, numpy.array_api is returned by this function.

This matches the behaviour I see in the python interpreter:

In [1]: import numpy.array_api as npa
<ipython-input-1-47d1a2ee1bbf>:1: UserWarning: The numpy.array_api submodule is still experimental. See NEP 47.
  import numpy.array_api as npa

In [2]: from scipy._lib.array_api_compat.array_api_compat import array_namespace

In [3]: x = npa.asarray([1, 2, 3, 4])

In [4]: xp = array_namespace(x)

In [5]: isinstance(x, npa._array_object.Array)
Out[5]: True

In [6]: xp
Out[6]: <module 'numpy.array_api' from '/home/lucas/mambaforge/envs/scipy-dev/lib/python3.10/site-packages/numpy/array_api/__init__.py'>

In [7]: xp.__name__
Out[7]: 'numpy.array_api'

[tupui]

from scipy._lib.array_api_compat.array_api_compat import array_namespace

Here you are actually calling array_api_compat.array_namespace. But we have our own:

from scipy._lib._array_api import array_namespace

---

EDIT: oh I think I see what you mean. Let me think about that.

[tupui]

Ok yes, so array_api_compat.array_namespace is indeed going to return numpy.array_api if you pass in numpy.array_api. This is a pass through to allow to check the standard. So there is no wrapping. Internally, the compat library look for __array_namespace__ and as it founds it for numpy.array_api, it directly returns this namespace.

[lucascolley]

Yep that's it. Matt, was this with the SCIPY_ARRAY_API flag set or not?

The output is an ndarray even when the input is an array_api.Array. That's because xp.full and xp.empty create ndarrays.

[mdhaber]

Yeah, not set.

[mdhaber]

Suggested change

	assert_equal(np.asarray(res1), ref1)
	assert_equal(np.asarray(res2), ref2)
	assert_equal(np.asarray(res1), ref1)
	assert_equal(res1.shape, ref1.shape)
	assert_equal(np.asarray(res2), ref2)
	assert_equal(res2.shape, ref2.shape)

[lucascolley]

A set_assert_equal utility (like set_assert_allclose discussed in the special PR) could be useful here, to allow us to use the equivalent from the other libraries we want to test with.

[mdhaber]

I meant to add a comment in the code about that to preempt this comment. But let's avoid scope creep - when set_assert_equal exists, we can use it, but that's not part of this PR. Same with the isinstance check below.

[mdhaber]

That said, if you'd like to submit a PR for set_assert_equal (and perhaps set_assert_allclose separate from gh-19005) I would be willing to review it. Rather than all these separate functions, though, why not something like an array_testing_namespace?

[lucascolley]

Okay, I may be able to submit that PR in the next few days 👍 but please don't wait for me if this is ready to go in. In its current state, I think that you will need to add @skip_if_array_api_gpu, since np.asarray fails for CuPy arrays.

[mdhaber]

Sure.

[lucascolley]

@mdhaber would you mind elaborating a bit on how an array_testing_namespace would work? Do we want testing functions where we pass in xp as a keyword, or do we keep the set_...(xp) structure? Or something else?

[mdhaber]

As array_namespace accepts arrays and returns a namespace with the appropriate Array API functions, array_testing_namespace would accept arrays and return a namespace with the appropriate testing functions. As array_namespace uses array_api_compat to give a consistent signature to existing functions, so array_testing_namespace would give a consistent signature to the various existing testing functions.

There are really not many array comparison testing functions to be handled. I really only need one to assess whether arrays are close (assert_allclose), equal (assert_equal or assert_array_equal, which don't seem different enough to exist separately aside from backward compatibility reasons), or ordered (assert_array_less).

So on second thought, rather than accepting arrays and returning a namespace, it might be better to just write three functions (e.g. assert_equal, assert_close, and assert_less) that accept two arrays, each of any type, and do the comparison as appropriate to the devices the arrays live on, whatever that means. I assume that for two arrays on the GPU we would use CuPy's functions, two arrays accessible to the CPU would use NumPy or PyTorch test functions, and arrays split between the devices would move the data to one or the other before comparison.

[rgommers]

That sounds good to me - agreed that we only need those couple of functions.

[mdhaber]

I think this suggestion is resolved, but I'll leave it open for others to refer to.

[mdhaber]

Suggested change

	@array_api_compatible
	@skip_if_array_api_gpu
	@array_api_compatible

Consider adding per #18930 (comment) (with required import at the top).

Alternatively, since array API behavior is opt-in and CI doesn't use GPUs right now, leave as-is. A failure in the future will remind us to replace np.assert_equal below with a GPU-compatible version (which will presumably have been written by the time anyone sees a failure).

[mdhaber]

@lucascolley IIUC this test should now pass with CuPy arrays, so we can resolve this? Can you test locally?

[lucascolley]

Looks fine with CuPy. python dev.py test -t scipy/_lib/tests/test__util.py -b all gives me this error with torch:

scipy/_lib/tests/test__util.py:364: in test_basic
    @array_api_compatible
        f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f22e371e980>
        self       = <scipy._lib.tests.test__util.TestLazywhere object at 0x7f22e3d9bb50>
        xp         = <module 'torch' from '/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/torch/__init__.py'>
scipy/_lib/tests/test__util.py:399: in test_basic
    xp_assert_equal(res1, ref1)
        arrays     = [tensor(0.), tensor(nan)]
        cond       = tensor(True)
        cond_shape = ()
        data       = data(...)
        dtype      = <class 'numpy.float32'>
        f          = <function TestLazywhere.test_basic.<locals>.f at 0x7f2342606e80>
        f2         = <function TestLazywhere.test_basic.<locals>.f2 at 0x7f234263a3e0>
        fillvalue  = tensor(0.)
        input_shapes = ((), (), ())
        mbs        = mutually_broadcastable_shapes(num_shapes=3, min_side=0)
        n_arrays   = 2
        p          = 0.0
        ref1       = tensor(nan)
        ref2       = tensor(nan)
        res1       = tensor(nan)
        res2       = tensor(nan)
        result_shape = ()
        rng        = Generator(PCG64) at 0x7F234263E5E0
        rng_seed   = 1000000000
        self       = <scipy._lib.tests.test__util.TestLazywhere object at 0x7f22e3d9bb50>
        shapes     = [(), ()]
        xp         = <module 'torch' from '/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/torch/__init__.py'>
scipy/_lib/_array_api.py:188: in assert_equal
    return xp.testing.assert_close(actual, desired, rtol=0, atol=0,
E   AssertionError
E   Falsifying example: test_basic(
E       self=<scipy._lib.tests.test__util.TestLazywhere object at 0x7f22e3d9bb50>,
E       xp=<module 'torch' from '/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/torch/__init__.py'>,
E       n_arrays=2,
E       rng_seed=1000000000,
E       dtype=numpy.float32,
E       p=0.0,
E       data=data(...),
E   )
E   Draw 1: BroadcastableShapes(input_shapes=((), (), ()), result_shape=())
E   Draw 2: array(0., dtype=float32)
E   Draw 3: array(0., dtype=float32)
E   Draw 4: array(nan, dtype=float32)
E   
E   You can reproduce this example by temporarily adding @reproduce_failure('6.82.7', b'AXicY2RgZCANKDAy/P8BYQIAEN8CGw==') as a decorator on your test case
        actual     = tensor(nan)
        desired    = tensor(nan)
        err_msg    = ''
        xp         = <module 'scipy._lib.array_api_compat.array_api_compat.torch' from '/home/lucas/dev/myscipy/build-install/lib/python3.11/site-packages/scipy/_lib/array_api_compat/array_api_compat/torch/__init__.py'>
========================================================================================= short test summary info ==========================================================================================
FAILED scipy/_lib/tests/test__util.py::TestLazywhere::test_basic[torch] - AssertionError
====================================================================================== 1 failed, 24 passed in 26.03s =======================================================================================

[mdhaber]

Thanks. Can you try again after implementing https://github.com/scipy/scipy/pull/19186/files#r1316228638.

[lucascolley]

python dev.py test -t scipy/_lib/tests/test__util.py -b all passes all tests after that patch yup.

However, with SCIPY_DEVICE=cuda, I get this error:

____________________________________________________________________________________ TestLazywhere.test_basic[torch] ______________________________________________________________________________________
scipy/_lib/tests/test__util.py:364: in test_basic
    @array_api_compatible
        f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f9b0639fc40>
        self       = <scipy._lib.tests.test__util.TestLazywhere object at 0x7f9b066ba390>
        xp         = <module 'torch' from '/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/torch/__init__.py'>
scipy/_lib/tests/test__util.py:384: in test_basic
    res1 = _lazywhere(cond, arrays, f, fillvalue)
        arrays     = [tensor([])]
        cond       = tensor(True)
        cond_shape = ()
        data       = data(...)
        dtype      = <class 'numpy.float32'>
        f          = <function TestLazywhere.test_basic.<locals>.f at 0x7f9be61a6de0>
        f2         = <function TestLazywhere.test_basic.<locals>.f2 at 0x7f9b652c2020>
        fillvalue  = tensor(0.)
        input_shapes = ((), (0,))
        mbs        = mutually_broadcastable_shapes(num_shapes=2, min_side=0)
        n_arrays   = 1
        p          = 0.0
        result_shape = (0,)
        rng        = Generator(PCG64) at 0x7F9B652E5FC0
        rng_seed   = 1000000000
        self       = <scipy._lib.tests.test__util.TestLazywhere object at 0x7f9b066ba390>
        shapes     = [(0,)]
        xp         = <module 'torch' from '/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/torch/__init__.py'>
scipy/_lib/_util.py:119: in _lazywhere
    out[cond] = temp1
E   RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Since cond is created from a numpy array (cond = xp.asarray(rng.random(size=cond_shape) > p)), it is created on the cpu.

Edit: I'm trying to test whether pytorch/pytorch#106773 has fixed this when the input is a numpy array, but I'm running into some issues https://discuss.pytorch.org/t/cuda-not-available-in-nightly-but-is-for-2-0-1/187767.

[mdhaber]

... it is created on the cpu.

Do I understand that:

• cond = xp.asarray(rng.random(size=cond_shape) > p) generates a CPU array
• arrays = [xp.asarray(data.draw(npst.arrays(dtype=dtype, shape=shape))) for shape in shapes] also generates CPU arrays
• xp = array_namespace(cond, *arrays) returns a Pytorch namespace
• temp1 = xp.asarray(f(*(arr[cond] for arr in arrays))) returns a CPU array
• out = xp.full(cond.shape, fill_value=fillvalue, dtype=dtype) generates a GPU array because I didn't specify the device.
• out[cond] = temp1 fails

This sounds like it is going to be a common gotcha. I would not expect the namespace returned by array_namespace(cond, *arrays) to have functions that produce arrays that are incompatible with cond and arrays.

[lucascolley]

I think that may be the case, but I don't know how asarray is supposed to interact with the default device and whether pytorch/pytorch#106773 will apply to numpy array inputs too.

The ideal scenario is that xp.asarray transfers to GPU à la CuPy when we have SCIPY_DEVICE=cuda.

I'll let you know once I hear back from PyTorch and am able to test with the nightly.

[lucascolley]

@mdhaber crisis averted: the pytorch/pytorch#106773 fix applies to np inputs too, so with a recent pytorch nightly cond = xp.asarray(rng.random(size=cond_shape) > p) will generate a GPU array when we have SCIPY_DEVICE=cuda.

So everything in this PR is looking good, the errors will go away in the next pytorch release.

[mdhaber]

Cool. Thanks.

[mdhaber]

So everything in this PR is looking good

Thanks. In that case, please consider hitting the approve button. (As a maintainer, it's helpful to see that someone else has also taken a look.)

[mdhaber]

@lucascolley A few thoughts about the new _array_api assertions:

• Should they be named something different in case they need to be used in the same file as np.testing assertions? Here, I've done the import assert_equal as xp_assert_equal.
• Does passing the assertion already imply that the result and reference are the same array type? If not, should it?
• What do you think of adding options that check the shape and dtype of the resulting array? Although a little inconvenient if all elements of an array are supposed to be the same number (e.g. 0), I would even suggest that checking shape and dtype should be part of the default comparison behavior to force us to be more careful about these things as we head into the Array API world.

[lucascolley]

Should they be named something different in case they need to be used in the same file as np.testing assertions? Here, I've done the import assert_equal as xp_assert_equal.

I'm happy to add the xp_ prefix.

Does passing the assertion already imply that the result and reference are the same array type? If not, should it?

I don't think that this is implied - maybe it should be. However, testing that the correct namespace is returned should always be done at least somewhere other than these functions. For fft I have a separate class called TestNamespaces, but these could be worked into the existing tests to be slightly more efficient if needed. These tests use Tyler's _assert_matching_namespace.

What do you think of adding options that check the shape and dtype of the resulting array? Although a little inconvenient if all elements of an array are supposed to be the same number (e.g. 0), I would even suggest that checking shape and dtype should be part of the default comparison behavior to force us to be more careful about these things as we head into the Array API world.

Yep, this sounds good. I'm not sure whether they even need to be options, we should probably always do this. torch does this anyway, so I'll just need to change np and cupy.

I'll make a follow-up PR with these improvements. We can always come back and tweak again as needed.

[mdhaber]

However, testing that the correct namespace is returned should always be done at least somewhere other than these functions.

Why? That is, why can't it just be baked into all the tests that also check the numerical results?

[lucascolley]

why can't it just be baked into all the tests that also check the numerical results?

Some functions will be tested with these assertions, some won't. To me, having a separate class to test namespaces, or at least a separate line in the test functions with that explicit purpose, will ensure that we have full coverage of these checks.

For groups of functions which are all tested using these functions, having it baked in would be sufficient, but I don't really like how implicit that is.

What do you think?

[lucascolley]

I thought that python dev.py test -t scipy/_lib/tests/test__util.py -b all was passing for me yesterday, but now I am getting a failure with torch, where res1 and ref1 are both nan:

_____________________________________________________________________________________ TestLazywhere.test_basic[torch] ______________________________________________________________________________________
scipy/_lib/tests/test__util.py:364: in test_basic
    @pytest.mark.filterwarnings('ignore::RuntimeWarning')  # overflows, etc.
        f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7fc27cc78040>
        self       = <scipy._lib.tests.test__util.TestLazywhere object at 0x7fc27d02c2d0>
        xp         = <module 'torch' from '/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/torch/__init__.py'>
/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/hypothesis/core.py:790: in execute_once
    result = self.test_runner(data, run)
        data       = ConjectureData(VALID, 61 bytes)
        example_kwargs = None
        expected_failure = None
        is_final   = True
        print_example = True
        run        = <function StateForActualGivenExecution.execute_once.<locals>.run at 0x7fc27d00bba0>
        self       = <hypothesis.core.StateForActualGivenExecution object at 0x7fc27d02c4d0>
        test       = <function TestLazywhere.test_basic at 0x7fc27cc77e20>
        text_repr  = None
/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/hypothesis/executors.py:47: in default_new_style_executor
    return function(data)
        data       = ConjectureData(VALID, 61 bytes)
        function   = <function StateForActualGivenExecution.execute_once.<locals>.run at 0x7fc27d00bba0>
/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/hypothesis/core.py:786: in run
    return test(*args, **kwargs)
        a          = []
        args       = ()
        argslices  = {'dtype': (6, 7), 'n_arrays': (0, 1), 'p': (7, 18), 'rng_seed': (1, 6)}
        context    = <hypothesis.control.BuildContext object at 0x7fc3b545bb10>
        data       = ConjectureData(VALID, 61 bytes)
        example_kwargs = None
        expected_failure = None
        is_final   = True
        kw         = {'data': data(...), 'dtype': <class 'numpy.float32'>, 'n_arrays': 2, 'p': 0.0, ...}
        kwargs     = {'data': data(...), 'dtype': <class 'numpy.float32'>, 'n_arrays': 2, 'p': 0.0, ...}
        print_example = True
        printer    = <hypothesis.vendor.pretty.RepresentationPrinter object at 0x7fc27d077e10>
        self       = <hypothesis.core.StateForActualGivenExecution object at 0x7fc27d02c4d0>
        test       = <function TestLazywhere.test_basic at 0x7fc27cc77e20>
        text_repr  = None
scipy/_lib/tests/test__util.py:401: in test_basic
    xp_assert_equal(res1, ref1)
        arrays     = [tensor(0.), tensor(nan)]
        cond       = tensor(True)
        cond_shape = ()
        data       = data(...)
        dtype      = <class 'numpy.float32'>
        f          = <function TestLazywhere.test_basic.<locals>.f at 0x7fc3b7f8a200>
        f2         = <function TestLazywhere.test_basic.<locals>.f2 at 0x7fc27cc15760>
        fillvalue  = tensor(0.)
        input_shapes = ((), (), ())
        mbs        = mutually_broadcastable_shapes(num_shapes=3, min_side=0)
        n_arrays   = 2
        p          = 0.0
        ref1       = tensor(nan)
        ref2       = tensor(nan)
        res1       = tensor(nan)
        res2       = tensor(nan)
        result_shape = ()
        rng        = Generator(PCG64) at 0x7FC27CC70200
        rng_seed   = 1000000000
        self       = <scipy._lib.tests.test__util.TestLazywhere object at 0x7fc27d02c2d0>
        shapes     = [(), ()]
        xp         = <module 'torch' from '/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/torch/__init__.py'>
scipy/_lib/_array_api.py:188: in assert_equal
    return xp.testing.assert_close(actual, desired, rtol=0, atol=0,
E   AssertionError
        actual     = tensor(nan)
        desired    = tensor(nan)
        err_msg    = ''
        xp         = <module 'scipy._lib.array_api_compat.array_api_compat.torch' from '/home/lucas/dev/myscipy/build-install/lib/python3.11/site-packages/scipy/_lib/array_api_compat/array_api_compat/torch/__init__.py'>
------------------------------------------------------------------------------------------------ Hypothesis ------------------------------------------------------------------------------------------------
Falsifying example: test_basic(
    self=<scipy._lib.tests.test__util.TestLazywhere object at 0x7fc27d02c2d0>,
    xp=<module 'torch' from '/home/lucas/mambaforge/envs/scipy-dev/lib/python3.11/site-packages/torch/__init__.py'>,
    n_arrays=2,
    rng_seed=1000000000,
    dtype=numpy.float32,
    p=0.0,
    data=data(...),
)
Draw 1: BroadcastableShapes(input_shapes=((), (), ()), result_shape=())
Draw 2: array(0., dtype=float32)
Draw 3: array(0., dtype=float32)
Draw 4: array(nan, dtype=float32)
========================================================================================= short test summary info ==========================================================================================
FAILED scipy/_lib/tests/test__util.py::TestLazywhere::test_basic[torch] - AssertionError
====================================================================================== 1 failed, 21 passed in 13.53s =======================================================================================

@reproduce_failure('6.82.7', b'AXicY2RgZCANKDAy/P8BYQIAEN8CGw==')

[mdhaber]

@lucascolley Isn't that the same as #18930 (comment), fixed by adding equal_nan=True in gh-19186?

[mdhaber]

What do you think?

These sort of separate checks (I'm lumping this in with checks of shape, dtype, and scalar vs array) have been forgotten for the past 20 years, so I imagine they will continue to be forgotten. But it seems to me that performing them automatically in these very widely used assertions can only help. I mean, I can anticipate an argument like "If we make it standard practice to include a separate namespace check separately, then it will never be forgotten, whereas if it is not written separately when these assertions are used, it may be forgotten when these assertions are not used." But I don't believe that statement.

[lucascolley]

Isn't that the same as #18930 (comment), fixed by adding equal_nan=True

Ah yes exactly, apologies, I forgot that I tested with that patch yesterday. I'll approve now 👍.

"If we make it standard practice to include a separate namespace check separately, then it will never be forgotten, whereas if it is not written separately when these assertions are used, it may be forgotten when these assertions are not used."

Yup, that's the argument I was thinking about.

But I don't believe that statement.

Fair enough. Happy to go along with whatever you think is best here 👍.

[lucascolley]

This all looks good to me now. python dev.py test -t scipy/_lib/tests/test__util.py -b all and the same with SCIPY_DEVICE=cuda should pass once #19186 is in and we get a new PyTorch release.

[tupui]

Thanks both. Looks like this is ready then. Let's get it in.