zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-31 16:21:27 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Power VJP fix for 0 (#505)

Browse Source
This commit is contained in:
Awni Hannun
2024-01-20 01:17:40 -08:00
committed by GitHub
parent 6bf779e72b
commit b207c2c86b
2 changed files with 19 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/tests/test_autograd.py
View File

@@ -380,6 +380,19 @@ class TestAutograd(mlx_tests.MLXTestCase):
out = mx.grad(fun)(mx.array(1.0, t), mx.array(1.0, t))
self.assertEqual(out.dtype, t)
def test_power_grad(self):
x = mx.array(0.0)
g = mx.grad(lambda x: x**2)(x)
self.assertEqual(g.item(), 0.0)
x = mx.array(0.0)
g = mx.grad(lambda x: x**1.5)(x)
self.assertEqual(g.item(), 0.0)
x = mx.array(2.0)
g = mx.grad(lambda x: x**2)(x)
self.assertAlmostEqual(g.item(), 4.0)
if __name__ == "__main__":
unittest.main()