fix readme

zinccat · Oct 22, 2024 · 9487828 · 9487828
1 parent 62298b5
commit 9487828
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -2,7 +2,9 @@
 
 Porting [FlexAttention](https://github.com/pytorch-labs/attention-gym) to pure JAX.
 
-Example usage:
+Please install Jax nightly: pip install -U --pre jax jaxlib "jax-cuda12-plugin[with_cuda]" jax-cuda12-pjrt -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
+
+Example usage: 
 
 ```python
 import jax

diff --git a/pyproject.toml b/pyproject.toml
@@ -12,6 +12,14 @@ homepage = "https://github.com/zinccat/flaxattention"
 python = "^3.10"
 jax = "^0.4.34"
 
+[tool.poetry.extras]
+viz = [
+    "matplotlib",
+    "numpy",
+    "jsonargparse",
+    "docstring-parser"
+]
+
 [tool.mypy]
 disable_error_code = "import-untyped"