zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 01:17:26 +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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 16 additions and 1 deletions
Show Stats 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(

15
python/tests/test_autograd.py
View File

@ -605,6 +605,21 @@ class TestAutograd(mlx_tests.MLXTestCase):
dfdx = mx.grad(fun)(x)
self.assertTrue(mx.allclose(dfdx, -2j * mx.ones_like(x)))
def test_flatten_unflatten_vjps(self):
def fun(x):
y = mx.unflatten(x, 0, (2, 2))
return y.sum()
x = mx.zeros((4, 8))
self.assertEqual(mx.grad(fun)(x).shape, (4, 8))
def fun(x):
y = mx.flatten(x, 0, 2)
return y.sum()
x = mx.zeros((2, 4, 8))
self.assertEqual(mx.grad(fun)(x).shape, (2, 4, 8))
if __name__ == "__main__":
unittest.main()