Fix scipy.special.logsumexp type promotion, test against NumPy and …

…SciPy nightlies (#643)

This PR enables testing against NumPy and SciPy nightlies, which were
added in #632 and later disabled to be split out here. I've also fixed a
test relating to incorrect shapes being cast (scalars were not being
treated as arrays) within the `combo_check` function, so all tests
should now pass on nightly and stable versions.

Related to #630

.github/workflows/test.yml

@@ -21,8 +21,7 @@ concurrency:
 
 jobs:
   test:
-    # name: Test / ${{ matrix.platform }} / Nightly ${{ matrix.nightly[0] }} / Python ${{ matrix.python-version }}
-    name: Test / ${{ matrix.platform }} / Python ${{ matrix.python-version }}
+    name: Regular tests / ${{ matrix.platform }} / Python ${{ matrix.python-version }}
     runs-on: ${{ matrix.platform }}
     strategy:
       fail-fast: false
@@ -40,12 +39,36 @@ jobs:
           allow-prereleases: true
       - uses: yezz123/setup-uv@v4
 
-      - name: Run CPython tests
-        if: ${{ !startsWith(matrix.python-version, 'pypy') }}
-        # run: uvx nox -s ${{ matrix.nightly[1] }}tests
+      # On PyPy, we skip SciPy because we don't have wheels
+      # available, see noxfile.py for more details.
+      - name: Run tests
         run: uvx nox -s tests
 
-      - name: Run PyPy tests
-        if: ${{ startsWith(matrix.python-version, 'pypy') }}
-        # run: uvx nox -s ${{ matrix.nightly[1] }}tests
-        run: uvx nox -s tests
+  # In this job, we test against the NumPy nightly wheels hosted on
+  # https://anaconda.org/scientific-python-nightly-wheels/numpy
+  # on the latest Python version available across platforms, instead of
+  # testing all Python versions and implementations on all platforms.
+  # We do not test on PyPy.
+  #
+  # However, "nox -s nightly-tests" can be used locally anywhere, on
+  # any Python version and implementation on any platform and we leave
+  # it to the user to decide what Python version to test against, which
+  # might or might not have a corresponding NumPy nightly wheel present.
+  nightlies:
+    name: Nightly tests / ${{ matrix.platform }} / Python ${{ matrix.python-version }}
+    runs-on: ${{ matrix.platform }}
+    strategy:
+      fail-fast: false
+      matrix:
+        platform: [ubuntu-latest, macos-13, macos-latest, windows-latest]
+        python-version: ["3.x"]
+
+    steps:
+      - uses: actions/[email protected]
+      - uses: actions/[email protected]
+        with:
+          python-version: ${{ matrix.python-version }}
+          allow-prereleases: true
+      - uses: yezz123/setup-uv@v4
+      - name: Run tests against nightly wheels for NumPy and SciPy
+        run: uvx nox -s nightly-tests

noxfile.py

@@ -50,7 +50,11 @@ def run_nightly_tests(session):
     session.install("-e", ".[test]", silent=False)
     # SciPy doesn't have wheels on PyPy
     if platform.python_implementation() == "PyPy":
-        session.install("numpy", "--upgrade", silent=False, env=UV_NIGHTLY_ENV_VARS)
+        session.install(
+            "numpy", "--upgrade", "--only-binary", ":all:", silent=False, env=UV_NIGHTLY_ENV_VARS
+        )
     else:
-        session.install("numpy", "scipy", "--upgrade", silent=False, env=UV_NIGHTLY_ENV_VARS)
+        session.install(
+            "numpy", "scipy", "--upgrade", "--only-binary", ":all:", silent=False, env=UV_NIGHTLY_ENV_VARS
+        )
     session.run("pytest", "--cov=autograd", "--cov-report=xml", "--cov-append", *session.posargs)

tests/test_scipy.py

@@ -228,7 +228,7 @@ def test_dirichlet_logpdf_alpha():
     ### Misc ###
     def test_logsumexp1():
         combo_check(special.logsumexp, [0], modes=["fwd", "rev"])(
-            [1.1, R(4), R(3, 4)], axis=[None, 0], keepdims=[True, False]
+            [np.array([1.1]), R(4), R(3, 4)], axis=[None, 0], keepdims=[True, False]
         )
 
     def test_logsumexp2():