google-research/timesfm · error · ValueError

Memory dimension ({self.qkv_features}) must be divisible by

Error message

Memory dimension ({self.qkv_features}) must be divisible by 'num_heads' heads ({self.num_heads}).

What it means

Multi-head attention splits the qkv/memory feature dimension into num_heads chunks of size head_dim, which requires qkv_features % num_heads == 0. The constructor validates this at build time and raises ValueError otherwise.

Source

Thrown at src/timesfm/flax/transformer.py:161

    use_rotary_position_embeddings: bool = True,
    use_bias: bool = False,
    deterministic: bool | None = None,
    attention_fn: Callable[..., Array] = nnx.dot_product_attention,
    qk_norm: str = "rms",
    rngs=nnx.Rngs(42),
  ):
    self.num_heads = num_heads
    self.in_features = in_features
    self.qkv_features = in_features
    self.out_features = in_features
    self.in_kv_features = in_features
    self.deterministic = deterministic
    self.use_bias = use_bias
    self.attention_fn = attention_fn
    self.qk_norm = qk_norm

    if self.qkv_features % self.num_heads != 0:
      raise ValueError(
        f"Memory dimension ({self.qkv_features}) must be divisible by "
        f"'num_heads' heads ({self.num_heads})."
      )
    self.head_dim = self.qkv_features // self.num_heads

    linear_general = functools.partial(
      LinearGeneral,
      out_features=(self.num_heads, self.head_dim),
      use_bias=self.use_bias,
    )
    # project inputs_q to multi-headed q/k/v
    # dimensions are then [batch..., length, n_heads, n_features_per_head]
    self.query = linear_general(self.in_features, rngs=rngs)
    self.key = linear_general(self.in_kv_features, rngs=rngs)
    self.value = linear_general(self.in_kv_features, rngs=rngs)

    if self.qk_norm == "rms":
      self.query_ln = RMSNorm(self.head_dim)

View on GitHub (pinned to 331c6d33cb)

Solutions

Choose num_heads that divides qkv_features evenly (e.g. model_dims=128 works with 1,2,4,8,16 heads).
Adjust qkv_features to the nearest multiple of num_heads if head count is fixed.
Add a config-level assertion/check early (at config load) to fail fast with a clear message.

Example fix

// before
attn = MultiHeadAttention(num_heads=7, qkv_features=128)  # ValueError
// after
attn = MultiHeadAttention(num_heads=8, qkv_features=128)  # head_dim = 16

Defensive patterns

Strategy: validation

Validate before calling

if config.model_dims % config.num_heads != 0:
    raise ValueError(f"model_dims ({config.model_dims}) must be divisible by num_heads ({config.num_heads})")

Try / catch

try:
    attn = MultiHeadAttention(num_heads=config.num_heads, qkv_features=config.model_dims)
except ValueError as e:
    if "divisible by" in str(e):
        config.num_heads = max(h for h in (1,2,4,8,16) if config.model_dims % h == 0)
        attn = MultiHeadAttention(num_heads=config.num_heads, qkv_features=config.model_dims)
    else:
        raise

Prevention

Choose num_heads from divisors of model_dims.
Validate the pairs (model_dims, num_heads) at config load.
Cross-check head counts when porting configs between model sizes.

When it happens

Trigger: Constructing MultiHeadAttention where qkv_features (often config.model_dims) is not divisible by num_heads, e.g. model_dims=100 with num_heads=4, or model_dims=128 with num_heads=7.

Common situations: Changing num_heads or model_dims independently in a config; adapting configs between model sizes where the two values were coupled; picking a head count that 'looks reasonable' without checking divisibility.

Related errors

AI-assisted analysis of google-research/timesfm@331c6d33cb (2026-08-29). Data as JSON: /api/errors/bbdcf212d1c19bf7. Report an issue: GitHub.