zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-18 01:50:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix unflatten vjp (#1708)

Browse Source
This commit is contained in:
Awni Hannun
2024-12-16 18:37:57 -08:00
committed by GitHub
parent a82996e9fb
commit d03c01dfbc
2 changed files with 16 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mlx/primitives.cpp
View File

@@ -1716,7 +1716,7 @@ std::vector<array> Unflatten::vjp(
const std::vector<array>& cotangents,
const std::vector<int>&,
const std::vector<array>&) {
return {flatten(cotangents[0], axis_, axis_ + shape_.size(), stream())};
return {flatten(cotangents[0], axis_, axis_ + shape_.size() - 1, stream())};
}
std::vector<array> Unflatten::jvp(