Fixes for fft/special modules (#14)

Author: lithomas1

scipy/_lib/meson.build

@@ -198,6 +198,7 @@ py3.install_sources(
     'array_api_compat/array_api_compat/numpy/_info.py',
     'array_api_compat/array_api_compat/numpy/fft.py',
     'array_api_compat/array_api_compat/numpy/linalg.py',
+    'array_api_compat/array_api_compat/numpy/_info.py',
   ],
   subdir: 'scipy/_lib/array_api_compat/numpy',
 )

scipy/fft/tests/test_helper.py

@@ -520,23 +520,24 @@ def test_definition(self, xp):
 
         # default dtype varies across backends
 
-        y = 9 * fft.fftfreq(9, xp=xp)
+        wrapped_xp = array_namespace(x)
+        y = 9 * fft.fftfreq(9, xp=wrapped_xp)
         xp_assert_close(y, x, check_dtype=False, check_namespace=True)
 
-        y = 9 * xp.pi * fft.fftfreq(9, xp.pi, xp=xp)
+        y = 9 * xp.pi * fft.fftfreq(9, xp.pi, xp=wrapped_xp)
         xp_assert_close(y, x, check_dtype=False)
 
-        y = 10 * fft.fftfreq(10, xp=xp)
+        y = 10 * fft.fftfreq(10, xp=wrapped_xp)
         xp_assert_close(y, x2, check_dtype=False)
 
-        y = 10 * xp.pi * fft.fftfreq(10, xp.pi, xp=xp)
+        y = 10 * xp.pi * fft.fftfreq(10, xp.pi, xp=wrapped_xp)
         xp_assert_close(y, x2, check_dtype=False)
 
     def test_device(self, xp):
         xp_test = array_namespace(xp.empty(0))
         devices = get_xp_devices(xp)
         for d in devices:
-            y = fft.fftfreq(9, xp=xp, device=d)
+            y = fft.fftfreq(9, xp=xp_test, device=d)
             x = xp_test.empty(0, device=d)
             assert xp_device(y) == xp_device(x)
 
@@ -552,23 +553,23 @@ def test_definition(self, xp):
         x2 = xp.asarray([0, 1, 2, 3, 4, 5], dtype=xp.float64)
 
         # default dtype varies across backends
-        
-        y = 9 * fft.rfftfreq(9, xp=xp)
+        wrapped_xp = array_namespace(x)
+        y = 9 * fft.rfftfreq(9, xp=wrapped_xp)
         xp_assert_close(y, x, check_dtype=False, check_namespace=True)
 
-        y = 9 * xp.pi * fft.rfftfreq(9, xp.pi, xp=xp)
+        y = 9 * xp.pi * fft.rfftfreq(9, xp.pi, xp=wrapped_xp)
         xp_assert_close(y, x, check_dtype=False)
 
-        y = 10 * fft.rfftfreq(10, xp=xp)
+        y = 10 * fft.rfftfreq(10, xp=wrapped_xp)
         xp_assert_close(y, x2, check_dtype=False)
 
-        y = 10 * xp.pi * fft.rfftfreq(10, xp.pi, xp=xp)
+        y = 10 * xp.pi * fft.rfftfreq(10, xp.pi, xp=wrapped_xp)
         xp_assert_close(y, x2, check_dtype=False)
 
     def test_device(self, xp):
         xp_test = array_namespace(xp.empty(0))
         devices = get_xp_devices(xp)
         for d in devices:
-            y = fft.rfftfreq(9, xp=xp, device=d)
+            y = fft.rfftfreq(9, xp=xp_test, device=d)
             x = xp_test.empty(0, device=d)
             assert xp_device(y) == xp_device(x)

scipy/special/_logsumexp.py

@@ -202,7 +202,10 @@ def _logsumexp(a, b, axis, return_sign, xp):
     a_max, i_max = _elements_and_indices_with_max_real(a, axis=axis, xp=xp)
 
     # for precision, these terms are separated out of the main sum.
-    a[i_max] = -xp.inf
+    # TODO: we shouldn't be mutating in-place here unless we make a copy
+    # dask arrays do not copy before this somehow
+    #a[i_max] = -xp.inf
+    a = xp.where(i_max, -xp.asarray(xp.inf, dtype=a.dtype), a)
     i_max_dt = xp.astype(i_max, a.dtype)
     # This is an inefficient way of getting `m` because it is the sum of a sparse
     # array; however, this is the simplest way I can think of to get the right shape.

scipy/special/tests/test_logsumexp.py

@@ -69,7 +69,11 @@ def test_logsumexp(self, xp):
         nan = xp.asarray([xp.nan])
         xp_assert_equal(logsumexp(inf), inf[0])
         xp_assert_equal(logsumexp(-inf), -inf[0])
-        xp_assert_equal(logsumexp(nan), nan[0])
+        # catch warnings here for dasks state there's no way to suppress
+        # warnings just for dask
+        # https://github.com/dask/dask/issues/3245
+        with np.errstate(divide='ignore', invalid='ignore'):
+            xp_assert_equal(logsumexp(nan), nan[0])
         xp_assert_equal(logsumexp(xp.asarray([-xp.inf, -xp.inf])), -inf[0])
 
         # Handling an array with different magnitudes on the axes