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

Pickle + dtype fix for numpy conversion (#763)

Browse Source 
* pickle + dtype fix for numpy conversion

* fix getattribute on Module base

* remove unused function

* fix tests

* add topk to ops

* fix doc
This commit is contained in:
Awni Hannun
2024-03-02 06:09:29 -08:00
committed by GitHub
parent 8e281c76c3
commit bc06cb9ff6
7 changed files with 99 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/tests/test_array.py
View File

@@ -1,6 +1,7 @@
# Copyright © 2023 Apple Inc.
import operator
import pickle
import unittest
import weakref
from itertools import permutations
@@ -1389,6 +1390,15 @@ class TestArray(mlx_tests.MLXTestCase):
b @= a
self.assertTrue(mx.array_equal(a, b))
def test_load_from_pickled_np(self):
a = np.array([1, 2, 3], dtype=np.int32)
b = pickle.loads(pickle.dumps(a))
self.assertTrue(mx.array_equal(mx.array(a), mx.array(b)))
a = np.array([1.0, 2.0, 3.0], dtype=np.float16)
b = pickle.loads(pickle.dumps(a))
self.assertTrue(mx.array_equal(mx.array(a), mx.array(b)))
if __name__ == "__main__":
unittest.main()