google-research/timesfm · error · ValueError

The embedding dims of the rotary position embeddingmust matc

Error message

The embedding dims of the rotary position embeddingmust match the hidden dimension of the inputs.

What it means

The rotary position embedding applies element-wise sinusoids that must align with the last dimension of the input tensor. If self.embedding_dims != inputs.shape[-1] the multiply would broadcast incorrectly, so forward() raises ValueError. This is an internal-shape invariant of the transformer.

Source

Thrown at src/timesfm/torch/transformer.py:77

  def __init__(
    self,
    embedding_dims: int,
    min_timescale: float = 1.0,
    max_timescale: float = 10000.0,
  ):
    super().__init__()
    self.embedding_dims = embedding_dims
    self.min_timescale = min_timescale
    self.max_timescale = max_timescale

  def forward(
    self,
    inputs: torch.Tensor,
    position: torch.Tensor | None = None,
  ):
    """Generates a JTensor of sinusoids with different frequencies."""
    if self.embedding_dims != inputs.shape[-1]:
      raise ValueError(
        "The embedding dims of the rotary position embedding"
        "must match the hidden dimension of the inputs."
      )
    half_embedding_dim = self.embedding_dims // 2
    fraction = (
      2
      * torch.arange(0, half_embedding_dim, device=inputs.device)
      / self.embedding_dims
    )
    timescale = (
      self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction
    ).to(inputs.device)
    if position is None:
      seq_length = inputs.shape[1]
      position = torch.arange(seq_length, dtype=torch.float32, device=inputs.device)[
        None, :
      ]

View on GitHub (pinned to 331c6d33cb)

Solutions

Ensure the input's last dimension equals the rotary embedding_dims (match TransformerConfig.model_dims).
If model_dims changed, rebuild the whole model/config so every layer uses the same dimension.
Apply RoPE before splitting into heads, or reshape back to (..., seq, model_dims) first; only use per-head RoPE if per-head dim equals embedding_dims.
Check any custom projection before RoPE: its output features must equal embedding_dims.

Example fix

// before
x = x.reshape(B, T, num_heads, head_dim)  # last dim = head_dim != model_dims
x = rotary(x)  # ValueError
// after
x = x.reshape(B, T, model_dims)
x = rotary(x)  # OK

Defensive patterns

Strategy: type-guard

Validate before calling

def can_apply_rotary(rotary, x):
    return x.shape[-1] == rotary.embedding_dims

if can_apply_rotary(rotary, x):
    x = rotary(x)

Type guard

def rotary_compatible(rotary, x) -> bool:
    return x.ndim in (3, 4) and x.shape[-1] == rotary.embedding_dims

Try / catch

try:
    x = rotary(x)
except ValueError:
    x = x.reshape(*x.shape[:-2], -1)  # merge heads back to model_dims
    x = rotary(x)

Prevention

Apply rotary embeddings before splitting into attention heads.
Keep model_dims consistent across all transformer layers and the rotary module.
Reshape tensors back to (..., seq, model_dims) before RoPE.
Add shape asserts in custom attention code paths.

When it happens

Trigger: Calling the rotary embedding forward with an input whose feature dimension differs from embedding_dims — typically when TransformerConfig.model_dims was changed inconsistently, or a reshaped/projected tensor with the wrong last-axis size is passed (e.g. per-head tensors whose last dim is head_dim).

Common situations: Modifying model_dims without rebuilding all layers consistently; feeding concatenated/split-head tensors to RoPE directly; custom attention code calling the RoPE module with transposed or reshaped inputs.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

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