Refactored CI and fixed types in tests/test_core.py

.github/workflows/ci.yml

@@ -21,4 +21,4 @@ jobs:
         run: rye test
 
       - name: Type-Check
-        run: rye run basedpyright -p . . 
+        run: rye run ci

pyproject.toml

@@ -3,7 +3,7 @@ name = "jaximal"
 version = "0.1.1"
 description = "A JAX-based PyTree manipulation library "
 authors = [
-    { name = "Arvin Kushwaha", email = "[email protected]" }
+    { name = "Arvin Kushwaha", email = "[email protected]" },
 ]
 dependencies = [
     "safetensors>=0.4.3",
@@ -22,12 +22,13 @@ build-backend = "hatchling.build"
 
 [tool.rye]
 managed = true
-dev-dependencies = [
-    "pytest>=8.2.1",
-    "basedpyright>=1.12.4",
-]
-excluded-dependencies = [
-]
+dev-dependencies = ["pytest>=8.2.1", "basedpyright>=1.12.4"]
+excluded-dependencies = []
+
+[tool.rye.scripts]
+ci = { chain = ["ci:verifytypes", "ci:basedpyright"] }
+"ci:verifytypes" = "rye run basedpyright --verifytypes jaximal"
+"ci:basedpyright" = "rye run basedpyright -p . ."
 
 [tool.hatch.metadata]
 allow-direct-references = true

tests/test_core.py

@@ -80,7 +80,7 @@ def forward(
     optimizer = optax.chain(optax.adam(1e-1), optax.contrib.reduce_on_plateau())
 
     mlp = MLP.init_state(3, 4, [16, 16], mlp_key)
-    opt_state: optax.OptState = optimizer.init(cast(optax.OptState, mlp))
+    opt_state: optax.OptState = optimizer.init(cast(optax.Params, mlp))
 
     def loss(
         mlp: MLP,
@@ -104,7 +104,7 @@ def update(
             value=cost,
         )
 
-        mlp: MLP = optax.apply_updates(mlp, updates)  # type: ignore
+        mlp = cast(MLP, optax.apply_updates(cast(optax.Params, mlp), updates))
 
         jax.debug.print('{} {}', i, cost)