huggingface/pytorch-image-models · error · ValueError

Unknown mode: {mode}

Error message

Unknown mode: {mode}

What it means

reshape_for_muon accepts only mode="flatten" or mode="batched" for reshaping >2-D conv weights; any other string raises this ValueError.

Source

Thrown at timm/optim/muon.py:365

    """
    original_shape = tensor.shape
    if tensor.ndim == 2:
        return tensor, original_shape
    if tensor.ndim < 2:
        raise ValueError(f"Tensor must have at least 2 dimensions, got {tensor.ndim}")

    out_ch, in_ch = tensor.shape[:2]
    if mode == "flatten":
        # Flatten: (out, in, *spatial) -> (out, in * spatial_prod)
        return tensor.reshape(out_ch, -1), original_shape
    elif mode == "batched":
        # Batched: (out, in, *spatial) -> (spatial_prod, out, in)
        # Move spatial dimension to front so zeropower_via_newtonschulz batches over it
        reshaped = tensor.reshape(out_ch, in_ch, -1)  # (out, in, spatial_prod)
        reshaped = reshaped.permute(2, 0, 1)  # (spatial_prod, out, in)
        return reshaped, original_shape
    else:
        raise ValueError(f"Unknown mode: {mode}")


def muon(
        params: List[torch.Tensor],
        grads: List[torch.Tensor],
        momentum_bufs: List[torch.Tensor],
        *,
        lr: float,
        weight_decay: float,
        momentum: float,
        nesterov: bool,
        ns_steps: int,
        ns_coefficients: NSCoeff,
        eps: float,
        safety_factor: float,
        adjust_lr_fn: Optional[str],
        conv_mode: str,
        normalize_spatial: bool,

View on GitHub (pinned to 9a5261e31b)

Solutions

  1. Use "flatten" or "batched" (these are also the valid conv_mode values of the Muon constructor)
  2. Set conv_mode on the Muon optimizer rather than calling the helper with an ad-hoc mode

Example fix

# before
reshape_for_muon(conv_weight, mode="flat")
# after
reshape_for_muon(conv_weight, mode="flatten")
Defensive patterns

Strategy: validation

Validate before calling

assert mode in ("flatten", "batched")

Type guard

def is_valid_mode(m: str) -> bool:
    return m in ("flatten", "batched")

Prevention

When it happens

Trigger: Calling reshape_for_muon(t, mode='reshape') or constructing Muon(conv_mode='flatten'/'batched' but passing a different mode string to the internal helper, or a typo like 'flat'.

Common situations: Typos in mode strings; code written against an older/newer API with different mode names.

Related errors


AI-assisted analysis of huggingface/pytorch-image-models@9a5261e31b (2026-08-27). Data as JSON: /api/errors/11d8e1a1a5b1e3f4. Report an issue: GitHub.