Adds prod and dev mode for tunix.

Author: wang2yn84

README.md

@@ -67,23 +67,16 @@ You can install Tunix in several ways:
 1. From PyPI (recommended):
 
 ```sh
-pip install tunix
+pip install "tunix[prod]"
 ```
 
-2. With development extras If you want to contribute or run tests and examples
-   locally:
-
-```
-pip install "tunix[dev]"
-```
-
-3. Directly from GitHub (latest main branch)
+2. Directly from GitHub (latest main branch)
 
 ```sh
 pip install git+https://github.com/google/tunix
 ```
 
-4. From source (editable install) If you plan to modify the codebase and run it
+3. From source (editable install) If you plan to modify the codebase and run it
    in development mode:
 
 ```sh

pyproject.toml

@@ -19,7 +19,6 @@ classifiers = [
 ]
 dependencies = [
   "datasets",
-  "flax>=0.12.0",
   "gcsfs",
   "grain",
   "huggingface_hub",
@@ -49,9 +48,11 @@ docs = [
     "sphinx-collections>=0.0.1",
     "sphinx_contributors",
 ]
+prod = [
+    "flax>=0.12.0",
+]
 dev = [
-    "flax@git+https://github.com/google/flax",
-    "qwix@git+https://github.com/google/qwix",
+    "flax>=0.11.2",
 ]
 
 [project.urls]