Merge pull request #375 from apax-hub/py12

Py 3.12 compatibility, rm chex and jaxtyping

Author: M-R-Schaefer

apax/md/md_checkpoint.py

@@ -1,13 +1,13 @@
 import logging
+import typing as t
 from pathlib import Path
 
 from flax.training import checkpoints
-from jaxtyping import PyTree
 
 log = logging.getLogger(__name__)
 
 
-def load_md_state(state: PyTree, ckpt_dir: Path) -> tuple[PyTree, int]:
+def load_md_state(state: t.Any, ckpt_dir: Path) -> tuple[t.Any, int]:
     try:
         log.info(f"loading MD state from {ckpt_dir}")
         target = {"state": state, "step": 0}

apax/nodes/optimizer/__init__.py

@@ -1,6 +1,6 @@
 from typing import NamedTuple
 
-import chex
+import jax
 import jax.numpy as jnp
 import optax
 from jax import tree_util as jtu
@@ -10,8 +10,8 @@
 
 
 class ScaleByAdemamixState(NamedTuple):
-    count: chex.Array
-    count_m2: chex.Array
+    count: jax.Array
+    count_m2: jax.Array
     m1: base.Updates
     m2: base.Updates
     nu: base.Updates