Merge pull request #2813 from AI-Hypercomputer:aireen/fix_pyconfig

PiperOrigin-RevId: 843338595

Author: Google-ML-Automation

src/MaxText/pyconfig.py

@@ -165,8 +165,11 @@ def tree_flatten(self):
 
   def __getattr__(self, attr: str) -> Any:
     """Provides attribute-style access to the final configuration dictionary."""
-    if attr in self._flat_config:
-      return self._flat_config[attr]
+    # Use object.__getattribute__ to avoid recursion when accessing _flat_config
+    # This is necessary for proper pickling/unpickling support
+    flat_config = object.__getattribute__(self, "_flat_config")
+    if attr in flat_config:
+      return flat_config[attr]
     raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{attr}'")
 
   def __setattr__(self, attr: str, value: Any) -> None: