Restore Python 3.9 compatibility

conftest.py

@@ -19,7 +19,7 @@ def pytest_addoption(parser):
 additional_modules += list((package_loc / "csdl_alpha" / "src" / "operations" / "special").glob("*.py"))
 additional_modules += list((package_loc / "csdl_alpha" / "src" / "operations" / "random").glob("*.py"))
 
-def pytest_collect_file(file_path, path, parent):
+def pytest_collect_file(file_path, parent):
     if file_path in additional_modules:
         return Module.from_parent(path=file_path, parent=parent)
     else:

csdl_alpha/backends/jax/utils.py

@@ -1,5 +1,24 @@
 from csdl_alpha.src.graph.operation import Operation
 from csdl_alpha.src.recorder import Recorder
+from functools import lru_cache
+
+@lru_cache(maxsize=None)
+def _pure_callback_has_vmap_method()->bool:
+    import inspect
+    import jax
+    return 'vmap_method' in inspect.signature(jax.pure_callback).parameters
+
+def sequential_pure_callback(callback, result_shape_dtypes, *args):
+    '''
+    Calls jax.pure_callback so that vmapped calls loop over the batch sequentially.
+
+    Newer JAX versions require vmap_method="sequential" for this, while older versions
+    (the latest supporting Python 3.9) don't accept vmap_method and are sequential by default.
+    '''
+    import jax
+    if _pure_callback_has_vmap_method():
+        return jax.pure_callback(callback, result_shape_dtypes, *args, vmap_method="sequential")
+    return jax.pure_callback(callback, result_shape_dtypes, *args)
 
 def fallback_to_inline_jax(operation:Operation, *args:list['jnp.array'])->tuple['jnp.array']:
     '''
@@ -12,11 +31,10 @@ def new_inline_func(*args_in):
         processed_inputs = [np.array(input) for input in args_in]
         return operation.compute_inline(*processed_inputs)
 
-    output = jax.pure_callback(
+    output = sequential_pure_callback(
         new_inline_func,
         [jax.ShapeDtypeStruct(output.shape, np.float64) for output in operation.outputs],
-        *args,
-        vmap_method="sequential")
+        *args)
     return tuple(output)
 
 

csdl_alpha/src/operations/custom/custom.py

@@ -83,6 +83,7 @@ def compute_inline(self, *args):
 
     def compute_jax(self, *args):
         import jax
+        from csdl_alpha.backends.jax.utils import sequential_pure_callback
 
         def new_inline_func(*args):
             processed_inputs = [np.array(input) for input in args]
@@ -94,11 +95,10 @@ def new_inline_func(*args):
         else:
             dtype = np.float32
 
-        output = jax.pure_callback(
+        output = sequential_pure_callback(
             new_inline_func,
             [jax.ShapeDtypeStruct(self.output_dict[output_var].shape, dtype) for output_var in self.output_dict],
-            *args,
-            vmap_method="sequential")
+            *args)
         # if len(output) == 1:
         #     output = output[0]
         return tuple(output)
@@ -450,6 +450,7 @@ def compute_inline(self, *orig_inputs_and_outputs_and_cots:list[np.array])->list
 
     def compute_jax(self, *args):
         import jax
+        from csdl_alpha.backends.jax.utils import sequential_pure_callback
 
         def new_inline_func(*args):
             processed_inputs = [np.array(input) for input in args]
@@ -461,11 +462,10 @@ def new_inline_func(*args):
         else:
             dtype = np.float32
 
-        output = jax.pure_callback(
+        output = sequential_pure_callback(
             new_inline_func,
             [jax.ShapeDtypeStruct(in_cot.shape, dtype) for in_cot in self.input_cotangents],
-            *args,
-            vmap_method="sequential")
+            *args)
         # if len(output) == 1:
         #     output = output[0]
 

csdl_alpha/src/operations/subop.py

@@ -36,13 +36,13 @@ def compute_inline(self, *args):
 
     def compute_jax(self, *args):
         from csdl_alpha.backends.jax import create_jax_function
+        from csdl_alpha.backends.jax.utils import sequential_pure_callback
         import jax
         jax_function = create_jax_function(self._subgraph, self.outputs, self.inputs)
         if self.jit:
-            output = jax.pure_callback(jax.jit(jax_function),
-                                       [jax.ShapeDtypeStruct(output.shape, np.float64) for output in self.outputs],
-                                       *args,
-                                       vmap_method="sequential")
+            output = sequential_pure_callback(jax.jit(jax_function),
+                                              [jax.ShapeDtypeStruct(output.shape, np.float64) for output in self.outputs],
+                                              *args)
             return tuple(output)
             # return tuple(jax.jit(jax_function)(*args))
         else:

docs/conf.py

@@ -38,7 +38,7 @@
     "sphinx_copybutton",            # allows copying code embedded in the docs rendered from .md or .ipynb files
     "myst_nb",                      # renders .md, .myst, .ipynb files
     "sphinx.ext.viewcode",          # adds the source code for classes and functions in auto generated api ref
-    "sphinxcontrib.collections",    # adds files from outside src and executes functions before Sphinx builds
+    "sphinx_collections",           # adds files from outside src and executes functions before Sphinx builds
     "sphinxcontrib.bibtex",         # for references and citations
 ]
 

requirements.txt

@@ -14,7 +14,7 @@ sphinx-autoapi==2.1.0
 astroid==2.15.5
 numpydoc
 gitpython==3.1.31
-git+https://github.com/anugrahjo/sphinx-collections.git
+sphinx-collections @ git+https://github.com/anugrahjo/sphinx-collections.git ; python_version >= "3.10"
 sphinxcontrib-bibtex==2.5.0
 rustworkx==0.14.2
 networkx