Make Kauldron modular when installing with pip.

PiperOrigin-RevId: 886101489

Author: ccrepy

README.md

@@ -27,6 +27,19 @@ Kauldron is a library for training machine learning models, optimized for
     post-mortem debugging on borg, runtime shape checking, and many others...).
 [Open an issue](https://github.com/google-research/kauldron/issues)..
 
+## Lightweight installation
+
+If you only need specific sub-modules, you can install them individually
+without pulling in the full dependency tree:
+
+```bash
+pip install kauldron[konfig]   # Config system only
+pip install kauldron[kontext]  # Tree manipulation / key system only
+pip install kauldron[ktyping]  # Array type checking only
+```
+
+A standard `pip install kauldron` still installs all dependencies.
+
 ## Citing Kauldron
 
 If Kauldron was helpful for a publication, please cite this repository:

pyproject.toml

@@ -13,9 +13,47 @@ classifiers = [
 ]
 keywords = []
 
-# pip dependencies of the project
-dependencies = [
+# By default, install all dependencies (backward compatible).
+# For lightweight installs, use e.g. `pip install kauldron[konfig]`
+dependencies = []
+
+# This is set automatically by flit using `kauldron.__version__`
+dynamic = ["version"]
+
+[project.urls]
+homepage = "https://github.com/google-research/kauldron"
+repository = "https://github.com/google-research/kauldron"
+documentation = "https://kauldron.readthedocs.io"
+
+# Other: `documentation`, `changelog`
+
+[project.optional-dependencies]
+# Lightweight sub-module installs
+konfig = [
     "absl-py",
+    "etils[edc,enp,epath,epy]",
+    "immutabledict",
+    "lark",
+    "ml_collections",
+]
+kontext = [
+    "etils[epy,etree-jax]",
+    "jax",
+    "ml_collections",
+    "typing_extensions",
+]
+ktyping = [
+    "etils[edc,enp,epath,epy]",
+    "immutabledict",
+    "jax",
+    "lark",
+    "numpy",
+    "typeguard>=4.4.1",
+]
+# Full install (all dependencies, equivalent to the previous default)
+# Installed through `pip install kauldron[all]`
+all = [
+    "kauldron[konfig,kontext,ktyping]",
     "altair",
     "chex",
     "clu",
@@ -24,13 +62,8 @@ dependencies = [
     "flax",
     "grain ; sys_platform != 'win32'",
     "graphviz",
-    "immutabledict",
-    "jax",
     "jaxtyping",
-    "lark",
     "mediapy",
-    "ml_collections",
-    "numpy",
     "opencv-python",  # speeds up pygrain image loading
     "optax",
     "orbax-checkpoint",
@@ -42,25 +75,12 @@ dependencies = [
     "tensorflow",
     "tensorflow_datasets>=4.9.7",
     "tqdm",
-    "typeguard>=4.4.1",
-    "typing_extensions",
     "xmanager",
 ]
-
-# This is set automatically by flit using `kauldron.__version__`
-dynamic = ["version"]
-
-[project.urls]
-homepage = "https://github.com/google-research/kauldron"
-repository = "https://github.com/google-research/kauldron"
-documentation = "https://kauldron.readthedocs.io"
-
-# Other: `documentation`, `changelog`
-
-[project.optional-dependencies]
 # Development deps (unittest, linting, formating,...)
 # Installed through `pip install .[dev]`
 dev = [
+    "kauldron[all]",
     "pytest",
     "pytest-xdist",
     "hypothesis==6.80.0",