Fix max iters issue and add tests

Author: goanpeca

.github/workflows/code-style.yml

@@ -75,6 +75,12 @@ jobs:
           conda activate $CONDA_ENV
           pip install uv
 
+      - name: Install GPG for Codecov
+        shell: bash -l {0}
+        run: |
+          # Install GPG which is required by codecov-action@v5
+          brew install gnupg
+
       - name: Install PyTorch
         if: ${{ matrix.pytorch-channel == 'pytorch' }}
         shell: bash -l {0}
@@ -130,7 +136,7 @@ jobs:
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v5
         with:
-          file: ${{ github.repository }}/coverage.xml
+          files: ${{ github.repository }}/coverage.xml
           flags: mps
           fail_ci_if_error: false
 

.github/workflows/mps-tests.yml

@@ -1,11 +1,18 @@
 from collections import OrderedDict
 from collections.abc import Mapping
-from typing import Tuple
+from typing import List, Tuple
 
 
 class Serializable:
-    _state_dict_all_req_keys: Tuple = ()
-    _state_dict_one_of_opt_keys: Tuple = ()
+    _state_dict_all_req_keys: Tuple[str, ...] = ()
+    _state_dict_one_of_opt_keys: Tuple[Tuple[str, ...], ...] = ((),)
+
+    def __init__(self) -> None:
+        self._state_dict_user_keys: List[str] = []
+
+    @property
+    def state_dict_user_keys(self) -> List:
+        return self._state_dict_user_keys
 
     def state_dict(self) -> OrderedDict:
         raise NotImplementedError
@@ -19,6 +26,21 @@ def load_state_dict(self, state_dict: Mapping) -> None:
                 raise ValueError(
                     f"Required state attribute '{k}' is absent in provided state_dict '{state_dict.keys()}'"
                 )
-        opts = [k in state_dict for k in self._state_dict_one_of_opt_keys]
-        if len(opts) > 0 and ((not any(opts)) or (all(opts))):
-            raise ValueError(f"state_dict should contain only one of '{self._state_dict_one_of_opt_keys}' keys")
+
+        # Handle groups of one-of optional keys
+        for one_of_opt_keys in self._state_dict_one_of_opt_keys:
+            if len(one_of_opt_keys) > 0:
+                opts = [k in state_dict for k in one_of_opt_keys]
+                num_present = sum(opts)
+                if num_present == 0:
+                    raise ValueError(f"state_dict should contain at least one of '{one_of_opt_keys}' keys")
+                if num_present > 1:
+                    raise ValueError(f"state_dict should contain only one of '{one_of_opt_keys}' keys")
+
+        # Check user keys
+        if hasattr(self, "_state_dict_user_keys") and isinstance(self._state_dict_user_keys, list):
+            for k in self._state_dict_user_keys:
+                if k not in state_dict:
+                    raise ValueError(
+                        f"Required user state attribute '{k}' is absent in provided state_dict '{state_dict.keys()}'"
+                    )