Method test_update_modules

python/tests/test_nn.py:295–331

(self)

            m.update(updates)

    def test_update_modules(self):
        m = nn.Sequential(nn.Linear(3, 3), nn.Linear(3, 3))

        # Updating non-existent modules should not be allowed by default
        with self.assertRaises(ValueError):
            m = m.update_modules({"values": [0, 1]})

        # Update wrong types
        with self.assertRaises(ValueError):
            m = m.update_modules({"layers": [0, 1]})

        class MyModule(nn.Module):
            def __init__(self):
                super().__init__()
                self.test = mx.array(1.0)
                self.list = [mx.array(1.0), mx.array(2.0)]

        m = MyModule()
        with self.assertRaises(ValueError):
            m = m.update_modules({"test": "hi"})
        with self.assertRaises(ValueError):
            m = m.update_modules({"list": ["hi"]})

        # Allow updating a strict subset
        m = nn.Sequential(nn.Linear(3, 3), nn.Linear(3, 3))
        m.update_modules({"layers": [{}, nn.Linear(3, 4)]})
        self.assertEqual(m.layers[1].weight.shape, (4, 3))

        # Using leaf_modules in the update should always work
        class MyModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.stuff = [nn.Linear(2, 2), 0, nn.Linear(2, 2)]
                self.more_stuff = {"hi": nn.Linear(2, 2), "bye": 0}

        m = MyModel()
        m.update_modules(m.leaf_modules())

    def test_parameter_deletion(self):
        m = nn.Linear(32, 32)