{"record":{"id":"0a8a2008a0d59fc5","repo":"google-research/timesfm","slug":"the-embedding-dims-of-the-rotary-position-embeddin","errorCode":null,"errorMessage":"The embedding dims of the rotary position embeddingmust match the hidden dimension of the inputs.","messagePattern":"The embedding dims of the rotary position embeddingmust match the hidden dimension of the inputs\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"src/timesfm/flax/transformer.py","lineNumber":87,"sourceCode":"\n  def __init__(\n    self,\n    embedding_dims: int,\n    min_timescale: int = 1,\n    max_timescale: int = 10000,\n  ):\n    self.embedding_dims = embedding_dims\n    self.min_timescale = min_timescale\n    self.max_timescale = max_timescale\n\n  def __call__(\n    self,\n    inputs: Float[Array, \"b ... d\"],\n    position: Array | None = None,\n  ):\n    \"\"\"Generates a JTensor of sinusoids with different frequencies.\"\"\"\n    if self.embedding_dims != inputs.shape[-1]:\n      raise ValueError(\n        \"The embedding dims of the rotary position embedding\"\n        \"must match the hidden dimension of the inputs.\"\n      )\n    half_embedding_dim = self.embedding_dims // 2\n    fraction = 2 * jnp.arange(0, half_embedding_dim) / self.embedding_dims\n    timescale = (\n      self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction\n    )\n    if position is None:\n      seq_length = inputs.shape[1]\n      position = jnp.arange(seq_length, dtype=jnp.float32)[None, :]\n    if len(inputs.shape) == 4:\n      position = position[..., None, None]\n      timescale = timescale[None, None, None, :]\n    elif len(inputs.shape) == 3:\n      position = position[..., None]\n      timescale = timescale[None, None, :]\n    else:","sourceCodeStart":69,"sourceCodeEnd":105,"githubUrl":"https://github.com/google-research/timesfm/blob/331c6d33cb1ac2611de3056d0ac7164aab6301eb/src/timesfm/flax/transformer.py#L69-L105","documentation":"Rotary position embeddings rotate pairs of input channels, so the rotary module's embedding_dims must equal the last (feature) dimension of the inputs tensor. The check runs in __call__, so it only fires at forward time when the configured rotary dims and actual hidden size diverge.","triggerScenarios":"Calling the rotary embedding with inputs whose last axis does not match self.embedding_dims, e.g. model_dims changed in config but the rotary embedding was built with the old dims, or feeding a tensor with the wrong feature count.","commonSituations":"Mismatches after changing model_dims in a transformer config without rebuilding layers, loading a checkpoint with different head/hidden dims, or passing an intermediate tensor of the wrong width into the rotary module directly.","solutions":["Align the rotary embedding's embedding_dims with the input feature dimension (inputs.shape[-1]).","If model_dims changed, re-instantiate the transformer/rotary module with the new dims rather than reusing the old one.","Inspect inputs.shape[-1] at the call site and compare to the config's model_dims to find where the divergence originates."],"exampleFix":"// before\nrotary = RotaryEmbedding(embedding_dims=64)\ny = rotary(x)  # x.shape[-1] == 128 -> ValueError\n// after\nrotary = RotaryEmbedding(embedding_dims=x.shape[-1])  # 128\ny = rotary(x)","handlingStrategy":"validation","validationCode":"assert rotary.embedding_dims == inputs.shape[-1], (\n    f\"rotary dims {rotary.embedding_dims} != input feature dim {inputs.shape[-1]}\")","typeGuard":"def rotary_dims_match(rotary, inputs) -> bool:\n    d = getattr(inputs, \"shape\", (None,))[-1]\n    return d == rotary.embedding_dims","tryCatchPattern":"try:\n    out = rotary(x)\nexcept ValueError as e:\n    if \"rotary position embedding\" in str(e):\n        rotary = RotaryEmbedding(embedding_dims=x.shape[-1])\n        out = rotary(x)\n    else:\n        raise","preventionTips":["Build the rotary module from the same config field (model_dims) used to size the input layers.","Re-instantiate all modules whenever model_dims changes.","Assert shape consistency in a forward-pass smoke test."],"tags":["shape","validation","jax","rotary-embedding"],"backgroundTag":"dimension-mismatch","analyzedSha":"331c6d33cb1ac2611de3056d0ac7164aab6301eb","analyzedAt":"2026-08-29T01:04:23.138Z","schemaVersion":2},"datasetVersion":"2026-08-29T02:17:18.158Z"}