Skip to content

[BUG][iOS] Incorrect output from ConvTransposed1d on large tensors #344

@cnnn

Description

@cnnn

Describe the bug
On iOS devices, ConvTransposed1d produces incorrect output on large input tensors, comparing with a reference implementation. The same code running on macOS does not have this issue, and the output of ConvTransposed1d matches exactly with the reference implementation. This indicates something on iOS metal backend silently truncates the input.

Comparing with our reference implementation, the error starts when T (timeIn) is between 8000 and 9000.

ConvTransposed1d Size Threshold Bug Test
Device: Device(gpu, 0)
============================================================
[OK  ] T=1,000                   | Ratio: 100.0% | MaxDiff: 0.0000
[OK  ] T=3,000                   | Ratio: 100.0% | MaxDiff: 0.0000
[OK  ] T=5,000                   | Ratio: 100.0% | MaxDiff: 0.0000
[OK  ] T=8,000                   | Ratio: 100.0% | MaxDiff: 0.0000
[FAIL] T=9,000                   | Ratio:  90.9% | MaxDiff: 1.3752
[FAIL] T=10,000                  | Ratio:  81.9% | MaxDiff: 1.5046
[FAIL] T=12,640 (Decoder L3)     | Ratio:  64.8% | MaxDiff: 1.4393
[FAIL] T=15,840 (Decoder L3)     | Ratio:  51.6% | MaxDiff: 1.5219
[FAIL] T=50,560 (Decoder L4)     | Ratio:  56.8% | MaxDiff: 1.7092

To Reproduce

Include code snippet

let input = MLXRandom.normal([1, timeIn, 384])
let conv = ConvTransposed1d(
    inputChannels: 384,
    outputChannels: 192,
    kernelSize: 8,
    stride: 4
)
let output = conv(input)  // Incorrect when timeIn is somewhere between 8000 and 9000 on iOS, while on macOS it produces correct output even when timeIn is as large as 50000.

Expected behavior
ConvTransposed1d should support large input tensors since our reference implementation can also achieve it, or at least error out when hitting some iOS-specific limitation, rather than failing silently.

Desktop (please complete the following information):

  • OS Version: macOS 26.1 (MacBook Pro M3 Max)
  • Device: iPhone 17 Pro Max, iPad Pro M5 2TB
  • Version iOS 26.1

Additional context
This issue might be related to MLX issue ml-explore/mlx#1377

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions