Rename AbstractPartialLoader to InitTransform

PiperOrigin-RevId: 869688526

Author: Conchylicultor

gemma/gm/ckpts/_checkpoint.py

@@ -46,7 +46,7 @@
 
 # TODO(epot): Should be part of core Kauldron
 @dataclasses.dataclass(frozen=True)
-class LoadCheckpoint(kd.ckpts.AbstractPartialLoader):
+class LoadCheckpoint(kd.ckpts.InitTransform):
   """Loads weights from a Gemma checkpoint.
 
   Note: The checpoint only contains the Gemma transformer weights, not the

gemma/gm/ckpts/_lora.py

@@ -34,10 +34,10 @@
 
 
 @dataclasses.dataclass(frozen=True)
-class SkipLoRA(kd.ckpts.AbstractPartialLoader):
+class SkipLoRA(kd.ckpts.InitTransform):
   """Wraps a partial loader to not restore the LoRA weights."""
 
-  wrapped: kd.ckpts.AbstractPartialLoader
+  wrapped: kd.ckpts.InitTransform
 
   def transform(self, state: _StateT) -> _StateT:  # pytype: disable=signature-mismatch
     # Remove the LoRA weights from the params structure so it can be restored

gemma/gm/ckpts/_policy.py

@@ -23,7 +23,7 @@
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
-class AnchoredPolicyLoader(kd.ckpts.AbstractPartialLoader):
+class AnchoredPolicyLoader(kd.ckpts.InitTransform):
   """Loader for `gm.nn.AnchoredPolicy` models.
 
   Loaded load policy and anchor separately by providing
@@ -33,8 +33,8 @@ class AnchoredPolicyLoader(kd.ckpts.AbstractPartialLoader):
   modifying the rest of the state.
   """
 
-  policy: kd.ckpts.AbstractPartialLoader
-  anchor: kd.ckpts.AbstractPartialLoader | None = None
+  policy: kd.ckpts.InitTransform
+  anchor: kd.ckpts.InitTransform | None = None
 
   def transform(self, state: kd.train.TrainState) -> kd.train.TrainState:
     if set(state.params.keys()) != {'policy', 'anchor'}: