update for new eqx

src/levanter/compat/hf_checkpoints.py

@@ -31,12 +31,7 @@
 import haliax
 from haliax import Axis
 from haliax.partitioning import ResourceMapping
-from haliax.state_dict import (
-    ModuleWithStateDictSerialization,
-    from_torch_compatible_state_dict,
-    save_state_dict,
-    to_torch_compatible_state_dict,
-)
+from haliax.state_dict import from_torch_compatible_state_dict, save_state_dict, to_torch_compatible_state_dict
 
 from levanter.logging import silence_transformer_nag
 from levanter.models.asr_model import ASRMixin
@@ -133,7 +128,7 @@ def hf_checkpoint_converter(cls) -> "HFCheckpointConverter":
 MConfig = TypeVar("MConfig", bound=HFCompatConfig)
 
 
-class ModelWithHfSerializationMixin(Generic[MConfig], ModuleWithStateDictSerialization):
+class ModelWithHfSerializationMixin(Generic[MConfig]):
     def get_hf_config(self):
         return self.config.to_hf_config(self.Vocab.size)
 

src/levanter/models/gpt2.py

@@ -245,7 +245,7 @@ def __call__(self, x: NamedArray, mask: Optional[AttentionMask | NamedArray], la
         return x
 
 
-class Gpt2Transformer(ModuleWithStateDictSerialization, eqx.Module):
+class Gpt2Transformer(ModuleWithStateDictSerialization):
     config: Gpt2Config = eqx.static_field()
     blocks: Stacked[Gpt2Block]
     ln_f: hnn.LayerNorm

src/levanter/models/lm_model.py

@@ -85,7 +85,7 @@ def build(self, Vocab: Axis, *, key: PRNGKey) -> "LmT":
         return self.model_type.init(Vocab, self, key=key)  # type: ignore
 
 
-class LmHeadModel(Generic[LmConfigT], abc.ABC):
+class LmHeadModel(eqx.Module, Generic[LmConfigT]):
     """
     Superclass for models with a language modeling head.
     """