huggingface / huggingface/candle

Matmul after reshape crashes on Metal

Open
#2,737 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
21.1k
Forks
1.8k
Avg merge
16h 42m
Merged PRs (30d)
25

Description

Doing a matrix multiplication after `transpose()` calls / after `permute()` crashes on Metal backend.
Here's a reproduction of the bug:

```rust
use candle_core::{DType, Device, Tensor};

pub fn main() -> candle_core::Result<()> {
let dev = Device::Cpu;
let x = Tensor::zeros(&[1, 8, 52, 1250], DType::F32, &dev)?;
let y = Tensor::zeros(&[1, 8, 1250, 52], DType::F32, &dev)?;
let z = x.matmul(&y)?; // works fine
dbg!(z.shape());

let x = Tensor::zeros(&[1, 8, 1250, 52], DType::F32, &dev)?;
let y = Tensor::zeros(&[1, 8, 52, 1250], DType::F32, &dev)?;
let z = x.matmul(&y)?; // works fine
dbg!(z.shape());

let x = Tensor::zeros(&[1, 1250, 416], DType::F32, &dev)?;
let y = Tensor::zeros(&[1, 1250, 416], DType::F32, &dev)?;

let x = x.reshape((1, 1250, 8, 52))?.transpose(1, 2)?;
let y = y
.reshape((1, 1250, 8, 52))?
.transpose(1, 2)?
.transpose(2, 3)?;

let z = x.matmul(&y)?; // works fine
dbg!(z.shape());

let dev = Device::new_metal(0)?;
let x = Tensor::zeros(&[1, 8, 52, 1250], DType::F32, &dev)?;
let y = Tensor::zeros(&[1, 8, 1250, 52], DType::F32, &dev)?;
let z = x.matmul(&y)?; // works fine
dbg!(z.shape());

let x = Tensor::zeros(&[1, 8, 1250, 52], DType::F32, &dev)?;
let y = Tensor::zeros(&[1, 8, 52, 1250], DType::F32, &dev)?;
let z = x.matmul(&y)?; // works fine
dbg!(z.shape());

let x = Tensor::zeros(&[1, 1250, 416], DType::F32, &dev)?;
let y = Tensor::zeros(&[1, 1250, 416], DType::F32, &dev)?;

let x = x.reshape((1, 1250, 8, 52))?.transpose(1, 2)?;
let y = y
.reshape((1, 1250, 8, 52))?
.transpose(1, 2)?
.transpose(2, 3)?;
// also crashes if you do permute(0, 2, 3, 1)
// also crashes if you create the array with shape (1, 1250, 8, 52) without reshaping

let z = x.matmul(&y)?; // crashes
dbg!(z.shape());

Ok(())
}
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by running the reproduction with Device::new_metal(0) and compare it with the CPU path. Trace Tensor::matmul after reshape, transpose, and permute, focusing on the Metal backend behavior; done means the shown non-contiguous inputs no longer crash and the result matches the working CPU cases.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
backend, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.