add jax extra to type checking task

Author: pattonw

.github/workflows/ty.yaml

@@ -14,4 +14,4 @@ jobs:
         uses: astral-sh/setup-uv@v6
 
       - name: ty
-        run: uv run --extra dev ty check src tests
+        run: uv run --extra dev --extra jax ty check src tests