diff --git a/conftest.py b/conftest.py index 6990c1c..8407dda 100644 --- a/conftest.py +++ b/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: diff --git a/csdl_alpha/backends/jax/utils.py b/csdl_alpha/backends/jax/utils.py index eceb8a5..05c7ab0 100644 --- a/csdl_alpha/backends/jax/utils.py +++ b/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) diff --git a/csdl_alpha/src/operations/custom/custom.py b/csdl_alpha/src/operations/custom/custom.py index ccd579a..bcf17f5 100644 --- a/csdl_alpha/src/operations/custom/custom.py +++ b/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] diff --git a/csdl_alpha/src/operations/subop.py b/csdl_alpha/src/operations/subop.py index b01a614..a5eb7b2 100644 --- a/csdl_alpha/src/operations/subop.py +++ b/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: diff --git a/docs/conf.py b/docs/conf.py index facc9dd..f0102b8 100644 --- a/docs/conf.py +++ b/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 ] diff --git a/requirements.txt b/requirements.txt index a465944..aa6c624 100644 --- a/requirements.txt +++ b/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