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

Author: awni

python/mlx/nn/layers/base.py

@@ -112,6 +112,7 @@ def __setattr__(self, key: str, val: Any):
             self[key] = val
         else:
             super(Module, self).__setattr__(key, val)
+            self.pop(key, None)
 
     def load_weights(
         self,

python/tests/test_nn.py

@@ -67,6 +67,12 @@ def initialize(self):
         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):