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

pick up preivously set attributes (#905)

Browse Source
This commit is contained in:
Awni Hannun 2024-03-26 11:19:59 -07:00 committed by GitHub
parent 9948eddf11
commit 570f2bf29e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
python/mlx/nn/layers/base.py
View File

@ -138,6 +138,11 @@ class Module(dict):
def __setattr__(self, key: str, val: Any): def __setattr__(self, key: str, val: Any):
if isinstance(val, (mx.array, dict, list, tuple)): if isinstance(val, (mx.array, dict, list, tuple)):
# If attribute was previously set but not in the
# dictionary, delete it so we pick it up in future
# calls to __getattr__
if hasattr(self, key) and key not in self:
delattr(self, key)
self[key] = val self[key] = val
else: else:
super(Module, self).__setattr__(key, val) super(Module, self).__setattr__(key, val)

15
python/tests/test_nn.py
View File

@ -54,6 +54,21 @@ class TestBase(mlx_tests.MLXTestCase):
m.apply_to_modules(assert_training) m.apply_to_modules(assert_training)
def test_module_attributes(self):
class Model(nn.Module):
def __init__(self):
super().__init__()
self.val = None
self.initialize()
def initialize(self):
self.val = mx.array(1.0)
model = Model()
self.assertTrue(mx.array_equal(model.val, mx.array(1.0)))
def test_model_with_dict(self): def test_model_with_dict(self):
class DictModule(nn.Module): class DictModule(nn.Module):
def __init__(self): def __init__(self):