Extract is_expert_param() helper to consolidate expert key matching

The same normalize_fqn + component-level matching logic existed in both
default_is_muon() and _expand_expert_params(). Extract into a single
is_expert_param() function in core.py so the logic lives in one place.

[skip-build]

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

torch-ext/optimizer/core.py

@@ -91,16 +91,25 @@ def adjust_lr_for_muon(lr, param_shape):
     return adjusted_lr
 
 
+def is_expert_param(name, expert_keys):
+    """Check if a parameter name matches any expert key (component-level)."""
+    if not expert_keys:
+        return False
+    parts = normalize_fqn(name).split(".")
+    return any(key in parts for key in expert_keys)
+
+
 def default_is_muon(name, x, expert_keys=None):
     normalized = normalize_fqn(name)
     parts = normalized.split(".")
     skip_keys = ["embed_tokens", "lm_head", "tok_embeddings", "output"]
     if any(key in parts for key in skip_keys):
-        logger.info("[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d",
-                    normalized, name, x.ndim)
+        logger.info(
+            "[is_muon] %s (orig: %s): skip (matched skip_key), ndim=%d",
+            normalized, name, x.ndim)
         return False
     effective_ndim = x.ndim
-    is_expert = expert_keys and any(key in parts for key in expert_keys)
+    is_expert = is_expert_param(name, expert_keys)
     if is_expert:
         effective_ndim -= 1
     result = effective_ndim >= 2

torch-ext/optimizer/muon.py

@@ -11,7 +11,7 @@ from torch.profiler import record_function
 from .adamw import step_adamw
 from .async_utils import run_pipeline
 from .core import (_muon_state, adjust_lr_for_muon,
-                   get_default_muon_param_groups, normalize_fqn, update_g,
+                   get_default_muon_param_groups, is_expert_param, update_g,
                    update_p)
 from .distributed.utils import (_is_shard, construct_shard_mesh,
                                 get_slices_of_dtensor)
@@ -46,8 +46,7 @@ def _expand_expert_params(names, params, expert_keys):
     expanded_params = []
 
     for n, p in zip(names, params):
-        is_expert = expert_keys and any(key in normalize_fqn(n).split(".")
-                                        for key in expert_keys)
+        is_expert = is_expert_param(n, expert_keys)
         is_dtensor = isinstance(p.data, DTensor)
 
         if not is_expert: