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

Fix logsumexp edge case (#740)

Browse Source 
* fix logsumexp

* fix inf constant

* also fix power grad

* fix ternary dispatch
This commit is contained in:
Awni Hannun
2024-02-25 08:39:55 -08:00
committed by GitHub
parent ac02cf33bd
commit e6418781ab
12 changed files with 112 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/tests/test_load.py
View File

@@ -66,13 +66,15 @@ class TestLoad(mlx_tests.MLXTestCase):
def test_save_and_load_safetensors(self):
if not os.path.isdir(self.test_dir):
os.mkdir(self.test_dir)
test_file = os.path.join(self.test_dir, "test.safetensors")
with self.assertRaises(Exception):
mx.save_safetensors("test", {"a": mx.ones((4, 4))}, {"testing": 0})
mx.save_safetensors(test_file, {"a": mx.ones((4, 4))}, {"testing": 0})
mx.save_safetensors(
"test", {"test": mx.ones((2, 2))}, {"testing": "test", "format": "mlx"}
test_file, {"test": mx.ones((2, 2))}, {"testing": "test", "format": "mlx"}
)
res = mx.load("test.safetensors", return_metadata=True)
res = mx.load(test_file, return_metadata=True)
self.assertEqual(len(res), 2)
self.assertEqual(res[1], {"testing": "test", "format": "mlx"})