Debug jax > 0.4.32 failures

.github/workflows/build-ci.yml

@@ -1,76 +1,76 @@
-name: CI
+# name: CI
 
-on:
-  pull_request:
-  push:
-    branches:
-      - main
-  schedule:
-      - cron: "0 13 * * 1"  # Every Monday at 9AM EST
+# on:
+#   pull_request:
+#   push:
+#     branches:
+#       - main
+#   schedule:
+#       - cron: "0 13 * * 1"  # Every Monday at 9AM EST
 
-jobs:
-  pre-commit:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v4
+# jobs:
+#   pre-commit:
+#     runs-on: ubuntu-latest
+#     steps:
+#       - name: Checkout repository
+#         uses: actions/checkout@v4
 
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-          cache: "pip"
-          cache-dependency-path: pyproject.toml
+#       - name: Set up Python
+#         uses: actions/setup-python@v5
+#         with:
+#           python-version: "3.10"
+#           cache: "pip"
+#           cache-dependency-path: pyproject.toml
 
-      - name: Test pre-commit hooks
-        run: |
-          python -m pip install --upgrade pip
-          pip install pre-commit
-          pre-commit run -a
+#       - name: Test pre-commit hooks
+#         run: |
+#           python -m pip install --upgrade pip
+#           pip install pre-commit
+#           pre-commit run -a
 
-  validate-types-and-docstrings:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v4
+#   validate-types-and-docstrings:
+#     runs-on: ubuntu-latest
+#     steps:
+#       - name: Checkout repository
+#         uses: actions/checkout@v4
 
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-          cache: "pip"
-          cache-dependency-path: pyproject.toml
+#       - name: Set up Python
+#         uses: actions/setup-python@v5
+#         with:
+#           python-version: "3.10"
+#           cache: "pip"
+#           cache-dependency-path: pyproject.toml
 
-      - name: Setup environment
-        run: |
-            python -m pip install --upgrade pip
-            pip install ".[tests,dev]"
+#       - name: Setup environment
+#         run: |
+#             python -m pip install --upgrade pip
+#             pip install ".[tests,dev]"
 
-      - name: mypy type validation
-        run: mypy src
+#       - name: mypy type validation
+#         run: mypy src
 
-      - name: darglint docstring validation
-        run: darglint src --strictness=short --ignore-raise=ValueError
+#       - name: darglint docstring validation
+#         run: darglint src --strictness=short --ignore-raise=ValueError
 
-  tests:
-    runs-on: ubuntu-latest
-    timeout-minutes: 5
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v4
+#   tests:
+#     runs-on: ubuntu-latest
+#     timeout-minutes: 5
+#     steps:
+#       - name: Checkout repository
+#         uses: actions/checkout@v4
 
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-          cache: "pip"
-          cache-dependency-path: pyproject.toml
+#       - name: Set up Python
+#         uses: actions/setup-python@v5
+#         with:
+#           python-version: "3.10"
+#           cache: "pip"
+#           cache-dependency-path: pyproject.toml
 
-      - name: Setup environment
-        run: |
-            python -m pip install --upgrade pip
-            pip install ".[tests,dev]"
+#       - name: Setup environment
+#         run: |
+#             python -m pip install --upgrade pip
+#             pip install ".[tests,dev]"
 
-      - name: Run Python tests
-        run: |
-          pytest --cov=invrs_opt tests
+#       - name: Run Python tests
+#         run: |
+#           pytest --cov=invrs_opt tests

.github/workflows/debug.yml

@@ -0,0 +1,79 @@
+name: CI
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+  schedule:
+      - cron: "0 13 * * 1"  # Every Monday at 9AM EST
+
+jobs:
+  run_debug_0020:
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+          cache: "pip"
+          cache-dependency-path: pyproject.toml
+
+      - name: Setup environment
+        run: |
+            python -m pip install --upgrade pip
+            pip install ".[tests,dev]"
+
+      - name: Run Python tests
+        run: |
+          python tests/optimizers/run_debug.py --steps=20
+
+  run_debug_0100:
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+          cache: "pip"
+          cache-dependency-path: pyproject.toml
+
+      - name: Setup environment
+        run: |
+            python -m pip install --upgrade pip
+            pip install ".[tests,dev]"
+
+      - name: Run Python tests
+        run: |
+          python tests/optimizers/run_debug.py --steps=100
+
+  run_debug_1000:
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+          cache: "pip"
+          cache-dependency-path: pyproject.toml
+
+      - name: Setup environment
+        run: |
+            python -m pip install --upgrade pip
+            pip install ".[tests,dev]"
+
+      - name: Run Python tests
+        run: |
+          python tests/optimizers/run_debug.py --steps=1000

pyproject.toml

@@ -16,7 +16,7 @@ maintainers = [
 ]
 
 dependencies = [
-    "jax < 0.4.32",
+    "jax",
     "jaxlib",
     "numpy",
     "requests",

tests/optimizers/run_debug.py

@@ -0,0 +1,48 @@
+"""Defines tests for the `lbfgsb.lbfgsb` module.
+
+Copyright (c) 2023 The INVRS-IO authors.
+"""
+
+import argparse
+
+import jax
+import jax.numpy as jnp
+from totypes import types
+
+from invrs_opt.optimizers import lbfgsb
+
+jax.config.update("jax_enable_x64", True)
+
+
+def optimization_with_vmap(steps):
+    params = types.Density2DArray(array=jnp.ones((2, 3, 3)))
+    opt = lbfgsb.density_lbfgsb(beta=2)
+    state = jax.vmap(opt.init)(params)
+
+    @jax.jit
+    @jax.vmap
+    def step_fn(state):
+        params = opt.params(state)
+        dummy_value = jnp.array(1.0, dtype=float)
+        dummy_grad = jax.tree_util.tree_map(jnp.ones_like, params)
+        state = opt.update(
+            grad=dummy_grad, value=dummy_value, params=params, state=state
+        )
+        return state, dummy_value
+
+    for i in range(steps):
+        print(f"vmap step {i}", flush=True)
+        state, value = step_fn(state)
+
+
+parser = argparse.ArgumentParser(prog="debug", description="opt debugging")
+parser.add_argument(
+    "--steps",
+    type=int,
+    default=None,
+    help="Number of steps.",
+)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    optimization_with_vmap(steps=args.steps)