| 75 | |
| 76 | def test_model_with_dict(self): |
| 77 | class DictModule(nn.Module): |
| 78 | def __init__(self): |
| 79 | super().__init__() |
| 80 | self.weights = {"w1": mx.zeros((2, 2)), "w2": mx.ones((2, 2))} |
| 81 | |
| 82 | model = DictModule() |
| 83 | params = tree_flatten(model.parameters(), destination={}) |
no outgoing calls