apply_model_overrides should follow the base signature

PiperOrigin-RevId: 702894126
Change-Id: Ifc240e6dcd28b9b661ac791cb93940cf1467f89e

saxml/server/pax/BUILD

@@ -144,5 +144,6 @@ pytype_strict_library(
         "//third_party/py/praxis:pax_fiddle",
         "//third_party/py/praxis:py_utils",
         "//third_party/py/praxis:pytypes",
+        "//third_party/py/typing_extensions",
     ],
 )

saxml/server/pax/union_model.py

@@ -27,6 +27,7 @@
 from praxis import pytypes
 from saxml.server.pax import servable_model
 from saxml.server.pax import servable_model_params
+from typing_extensions import override
 
 PRNGKey = pytypes.PRNGKey
 NestedMap = py_utils.NestedMap
@@ -69,10 +70,14 @@ def task(self) -> pax_fiddle.Config[base_task.BaseTask]:
   def datasets(self) -> List[pax_fiddle.Config[base_input.BaseInput]]:
     raise NotImplementedError('should not be called')
 
+  @override
   @classmethod
-  def apply_model_overrides(cls, overrides: Dict[str, Any]) -> None:
+  def apply_model_overrides(
+      cls, overrides: Dict[str, Any]
+  ) -> type['UnionModelParams']:
     """Delays the model overrides until child creation."""
     cls.overrides = overrides
+    return cls
 
 
 class UnionModel(servable_model.ServableModel):