Merge pull request MilesCranmer#1112 from MilesCranmer/fix-ci

ci: fix upper bound python for jax compat

Author: MilesCranmer

.github/workflows/CI.yml

@@ -206,6 +206,34 @@ jobs:
       - name: "Run tests"
         run: python -m pysr test main,jax,torch
 
+  autodiff_backends:
+    name: Test autodiff backends
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash -l {0}
+    strategy:
+      matrix:
+        python-version: ['3.13']
+        julia-version: ['1']
+
+    steps:
+      - uses: actions/checkout@v5
+      - name: "Set up Julia"
+        uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.julia-version }}
+      - name: "Set up Python"
+        uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+      - name: "Install PySR and all dependencies"
+        run: |
+            python -m pip install --upgrade pip
+            pip install '.[dev]'
+      - name: "Run autodiff backend tests"
+        run: python -m pysr test autodiff
   wheel_test:
     name: Test from wheel
     runs-on: ubuntu-latest

pysr/export.py

@@ -83,8 +83,8 @@ def add_export_formats(
     )
 
     if output_jax_format:
-        exports["jax_format"] = jax_format
+        exports["jax_format"] = pd.Series(jax_format, index=output.index)
     if output_torch_format:
-        exports["torch_format"] = torch_format
+        exports["torch_format"] = pd.Series(torch_format, index=output.index)
 
     return exports