Merge pull request #235 from inseq-team/bug/jaxlib

Add optional jax group to enforce compatible jaxlib version.
inseq-team · Nov 15, 2023 · 9d3f5c7 · 9d3f5c7
2 parents 4d8f688 + 078efce
commit 9d3f5c7
Show file tree

Hide file tree

Showing 3 changed files with 138 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -71,6 +71,8 @@ After installation, you should be able to run `make fast-test` and `make lint` w
 
 - Installing `sentencepiece` requires various packages, install with `sudo apt-get install cmake build-essential pkg-config` or `brew install cmake gperftools pkg-config`.
 
+- Inseq does not work with older versions of `jaxlib`.
+Install a compatible version with ```poetry install --with jax```.
 </details>
 
 ## Example usage in Python

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -101,6 +101,13 @@ pytest = ">=7.2.0"
 pytest-cov = ">=4.0.0"
 ruff  = ">=0.1.2"
 
+[tool.poetry.group.jax]
+optional = true
+
+[tool.poetry.group.jax.dependencies]
+jax = ">=0.4.20"
+jaxlib = ">=0.4.20"
+
 [tool.poetry.extras]
 sklearn = ["scikit-learn", "joblib"]
 datasets = ["datasets"]