No public description

PiperOrigin-RevId: 959730466

Author: h-joo

vmoe/checkpoints/base.py

@@ -118,7 +118,7 @@ def _parse_step_and_suffix_or_error(filepath):
       sorted(steps_and_suffixes - incomplete_steps_and_suffixes,
              reverse=decreasing),
       lambda x: x[0]):
-    if not suffixes.difference(set(x[1] for x in group)):
+    if not suffixes.difference(set(x[1] for x in group)):  # pyrefly: ignore[missing-attribute]
       yield step
 
 

vmoe/train/optimizer.py

@@ -166,7 +166,7 @@ def frozen_fn(params: optax.Params) -> PyTree:
     flatparams = flax.traverse_util.flatten_dict(
         flax.serialization.to_state_dict(params), sep='/')
     output = {
-        key: search_true if pattern.search(key) else search_false
+        key: search_true if pattern.search(key) else search_false  # pyrefly: ignore[missing-attribute]
         for key, value in flatparams.items()
     }
     return flax.serialization.from_state_dict(