Add pyproject.toml and uv install support

README.md

@@ -36,7 +36,7 @@ We implement both linear and log-complexity recurrent models.
 | Long Short-Term Memory Unit | $O(n)$ | [[paper]](https://ieeexplore.ieee.org/abstract/document/6795963) | [[code]](memax/equinox/set_actions/lstm.py) |
 
 # Datasets
-We provide [datasets](memax/datasets) to test our recurrent models. 
+We provide [datasets](memax/datasets) to test our recurrent models.
 
 ### Sequential MNIST [[HuggingFace]](https://huggingface.co/datasets/ylecun/mnist) [[Code]](memax/datasets/sequential_mnist.py)
 > The recurrent model receives an MNIST image pixel by pixel, and must predict the digit class.
@@ -54,7 +54,7 @@ We provide [datasets](memax/datasets) to test our recurrent models.
 > **Sequence Lengths:** `[20, 100, 1_000]`
 
 # Getting Started
-Install `memax` using pip and git for your specific framework
+Install `memax` using pip for your specific framework:
 ```bash
 pip install "memax[equinox]"
 pip install "memax[flax]"
@@ -65,6 +65,15 @@ pip install "memax[train,equinox]"
 pip install "memax[train,flax]"
 ```
 
+Or install from source using `uv`:
+
+```bash
+uv sync --extra equinox # Only equinox
+uv sync --extra flax # Only flax
+uv sync --extra train --extra equinox
+uv sync --extra train --extra flax
+```
+
 ## Equinox Quickstart
 ```python
 from memax.equinox.train_utils import get_residual_memory_model
@@ -76,12 +85,12 @@ from memax.equinox.train_utils import add_batch_dim
 T, F = 5, 6 # time and feature dim
 
 model = get_residual_memory_model(
-    model_name="LRU", input=F, hidden=8, output=1, num_layers=2, 
+    model_name="LRU", input=F, hidden=8, output=1, num_layers=2,
     key=jax.random.key(0)
 )
 
 starts = jnp.array([True, False, False, True, False])
-xs = jnp.zeros((T, F)) 
+xs = jnp.zeros((T, F))
 hs, ys = filter_jit(model)(model.initialize_carry(), (xs, starts))
 last_h = filter_jit(model.latest_recurrent_state)(hs)
 
@@ -101,7 +110,7 @@ python run_linen_experiments.py # flax linen framework
 ```
 
 
-## Custom Architectures 
+## Custom Architectures
 memax uses the [`equinox`](https://github.com/patrick-kidger/equinox) neural network library. See [the semigroups directory](memax/equinox/semigroups) for fast recurrent models that utilize an associative scan. We also provide a beta [`flax.linen`](https://flax-linen.readthedocs.io/en/latest/) API. In this example, we focus on `equinox`.
 
 ```python
@@ -148,7 +157,7 @@ print(debug_shape(h))
 #     (5,) # Start carries for first layer
 #     (5, 16) # Recurrent states of second layer
 #     (5,)) # Start carries for second layer
-# 
+#
 # Do your prediction
 prediction = jax.nn.softmax(y)
 

pyproject.toml

@@ -1,2 +1,42 @@
+[project]
+name = "memax"
+version = "0.1.1"
+description = "Deep memory and sequence modeling in JAX"
+readme = "README.md"
+requires-python = ">=3.9"
+dependencies = [
+  "jax",
+  "jaxlib",
+  "jaxtyping",
+  "optax",
+  "beartype",
+]
+
+[project.optional-dependencies]
+equinox = ["equinox"]
+flax = ["flax"]
+train = [
+  "datasets",
+  "tqdm",
+  "pillow",
+  "wandb",
+]
+all = [
+  "equinox",
+  "flax",
+  "datasets",
+  "tqdm",
+  "pillow",
+  "wandb",
+]
+test = ["pytest"]
+
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["."]
+
 [tool.isort]
 profile = "black"

requirements.lock.txt

@@ -0,0 +1,81 @@
+absl-py==2.4.0
+aiofiles==25.1.0
+aiohappyeyeballs==2.6.2
+aiohttp==3.13.5
+aiosignal==1.4.0
+annotated-doc==0.0.4
+annotated-types==0.7.0
+anyio==4.13.0
+attrs==26.1.0
+beartype==0.22.9
+certifi==2026.5.20
+charset-normalizer==3.4.7
+click==8.4.1
+datasets==4.8.5
+dill==0.4.1
+equinox==0.13.8
+etils==1.14.0
+filelock==3.29.0
+flax==0.12.7
+frozenlist==1.8.0
+fsspec==2026.2.0
+gitdb==4.0.12
+gitpython==3.1.50
+h11==0.16.0
+hf-xet==1.5.0
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==1.17.0
+humanize==4.15.0
+idna==3.17
+iniconfig==2.3.0
+jax==0.10.1
+jaxlib==0.10.1
+jaxtyping==0.3.10
+markdown-it-py==4.2.0
+mdurl==0.1.2
+ml-dtypes==0.5.4
+msgpack==1.1.2
+multidict==6.7.1
+multiprocess==0.70.19
+numpy==2.4.6
+opt-einsum==3.4.0
+optax==0.2.8
+orbax-checkpoint==0.11.40
+packaging==26.2
+pandas==3.0.3
+pillow==12.2.0
+platformdirs==4.10.0
+pluggy==1.6.0
+prometheus-client==0.25.0
+propcache==0.5.2
+protobuf==7.35.0
+psutil==7.2.2
+pyarrow==24.0.0
+pydantic==2.13.4
+pydantic-core==2.46.4
+pygments==2.20.0
+pytest==9.0.3
+python-dateutil==2.9.0.post0
+pyyaml==6.0.3
+requests==2.34.2
+rich==15.0.0
+scipy==1.17.1
+sentry-sdk==2.61.0
+shellingham==1.5.4
+simplejson==4.1.1
+six==1.17.0
+smmap==5.0.3
+tensorstore==0.1.84
+tqdm==4.67.3
+treescope==0.1.10
+typer==0.25.1
+typing-extensions==4.15.0
+typing-inspection==0.4.2
+urllib3==2.7.0
+uvloop==0.22.1
+wadler-lindig==0.1.7
+wandb==0.27.0
+xxhash==3.7.0
+yarl==1.24.2
+zipp==4.1.0