CoreMLTools PyTorch converter fails on rotary-attention einsum (BeatThis model) with rank/perm mismatch

🐞Describing the bug

Converting the BeatThis model (rotary attention, ISMIR 2024) from PyTorch to CoreML fails during einsum lowering. The converter throws a rank/perm mismatch (perm length 3 != rank(x) 1). PyTorch inference works fine; this appears to be a coremltools converter bug for this einsum pattern.

Source: https://github.com/CPJKU/beat_this

Stack Trace

ERROR - converting 'einsum' op (located at: 'base/frontend/blocks/0/partial/attnF/x.13'):

Traceback (most recent call last):
  ...
  File ".../coremltools/converters/mil/frontend/_utils.py", line 573, in solve_generic_einsum
    parsed_vectors, vars = solve_diagonal_einsum(parsed_vectors, vars)
  File ".../coremltools/converters/mil/frontend/_utils.py", line 495, in solve_diagonal_einsum_one_step
    x = mb.transpose(x=x, perm=perm)
  File ".../coremltools/converters/mil/mil/ops/defs/iOS15/tensor_transformation.py", line 988, in type_inference
    raise ValueError(msg.format(len(perm), self.x.rank))
ValueError: perm should have the same length as rank(x): 3 != 1

To Reproduce

import coremltools as ct
import torch
from beat_this.model.beat_tracker import BeatThis
from beat_this.utils import replace_state_dict_key

checkpoint = torch.load("beat_this/beat_this-final0.ckpt", map_location="cpu", weights_only=True)
hparams = {
    k: v for k, v in checkpoint["hyper_parameters"].items()
    if k in set(BeatThis.__init__.__code__.co_varnames)
}
model = BeatThis(**hparams)
state_dict = replace_state_dict_key(checkpoint["state_dict"], "model.", "")
model.load_state_dict(state_dict)
model.eval()

example = torch.randn(1, 1500, 128)
traced = torch.jit.trace(model, example, check_trace=False)

ct.convert(
    traced,
    inputs=[ct.TensorType(name="mel_spectrogram", shape=example.shape)],
)

If the model conversion succeeds, but there is a numerical mismatch in predictions, please include the code used for comparisons.

System environment (please complete the following information):

coremltools version: 9.0
macOS 15.7.3 (24G407) (arm64)
torch 2.2.0

Additional context

This happens in the rotary‑attention einsum pattern used by BeatThis. ONNX export with opset ≥12 works, but legacy ONNX conversion fails due to Slice inputs.

Can you give us a standalone example, which doesn't require loading an external model?

Can you give us a standalone example, which doesn't require loading an external model?

import torch
import coremltools as ct

class EinsumRepro(torch.nn.Module):
    def forward(self, rotations, freq_ranges):
        return torch.einsum('..., f -> ... f', rotations, freq_ranges)

model = EinsumRepro().eval()

rotations = torch.randn(2, 3)
freq_ranges = torch.randn(3)

traced = torch.jit.trace(model, (rotations, freq_ranges), check_trace=False)

ct.convert(
    traced,
    inputs=[
        ct.TensorType(name='rotations', shape=rotations.shape),
        ct.TensorType(name='freq_ranges', shape=freq_ranges.shape),
    ],
)

The einsum expression is valid in PyTorch, but coremltools appears to be fragile when handling automatic ellipsis (...) expansion.
For temparary workaround, you can use something like

    idx = "abcd"[:x.ndim]
    torch.einsum(f"{idx},f->{idx}f", x, y)

Repro hits on main (HEAD f272aa3). The failure is in parse_einsum_equation: each . in ...,f->...f maps to int 0, so the first parsed vector becomes [0, 0, 0]. solve_diagonal_einsum then treats three "duplicate" indices over a rank-2 tensor and the downstream transpose fails. Ellipsis is never resolved against actual input ranks.

One approach: before parse_einsum_equation, rewrite each ... to fresh single-char symbols whose count equals rank(input) - (len(segment) - 3), and reuse those symbols in the output. Two test cases: ...,f->...f and ...ij,...jk->...ik.

Would you take a PR for that, or prefer the ellipsis handling inside parse_einsum_equation itself?

@eyupcanakman

@eyupcanakman - I would accept a PR which takes either approach as long as it solves this issue. Please be sure to include a unit test which passes with your fix but fails without it.

Fix einsum ellipsis lowering to MIL (apple#2706)

Navigation Menu

Repository navigation

CoreMLTools PyTorch converter fails on rotary-attention einsum (BeatThis model) with rank/perm mismatch #2644

🐞Describing the bug

Stack Trace

To Reproduce

System environment (please complete the following information):

Additional context

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

CoreMLTools PyTorch converter fails on rotary-attention einsum (BeatThis model) with rank/perm mismatch #2644

Description

🐞Describing the bug

Stack Trace

To Reproduce

System environment (please complete the following information):

Additional context

Activity

TobyRoseman commented on Jan 20, 2026

tillt commented on Jan 21, 2026

jayisaking commented on Feb 6, 2026

eyupcanakman commented on May 9, 2026

TobyRoseman commented on May 18, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions