✅: jax 0.6.0+

Signed-off-by: nstarman <[email protected]>

Author: nstarman

.github/workflows/ci.yml

@@ -150,7 +150,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        jax-version: ["0.7.0"]
+        jax-version: ["0.6.2", "0.7.0"]
 
     steps:
       - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0