Now provides PyTreeDef, and can detect PyTrees via issubclass(x, PyTree)

2026-09-09 11:24:55 +08:00 · 2023-06-01 10:56:00 -07:00
parent 10e1852b37
commit 6a64ef114e
7 changed files with 68 additions and 43 deletions
@@ -42,7 +42,7 @@ jobs:
      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
-          python -m pip install pytest wheel beartype equinox jaxlib cloudpickle
+          python -m pip install -r test/requirements.txt
          python -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu

      - name: Checks with pre-commit