fix module attribute set, reset, set (#1403)

2025-09-04 06:44:40 +08:00 · 2024-09-11 16:30:42 -07:00
parent 02efb310ca
commit 8b30acd7eb
2 changed files with 7 additions and 0 deletions
--- a/python/tests/test_nn.py
+++ b/python/tests/test_nn.py
@@ -67,6 +67,12 @@ class TestBase(mlx_tests.MLXTestCase):
        model = Model()
        self.assertTrue(mx.array_equal(model.val, mx.array(1.0)))

+        model.val = None
+        self.assertEqual(model.val, None)
+
+        model.val = mx.array([3])
+        self.assertEqual(model.val.item(), 3)
+
    def test_model_with_dict(self):
        class DictModule(nn.Module):
            def __init__(self):