chore: import upstream snapshot with attribution

tests/python/relax/test_frontend_nn_parameter_containers.py

@@ -0,0 +1,219 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Any
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.relax.frontend import nn
+
+
+class ParamContainerModule(nn.Module):
+    def __init__(self):
+        self.list_params = nn.ParameterList(
+            [
+                nn.Parameter((4,), "float32"),
+                nn.Parameter((4,), "float32"),
+            ]
+        )
+        self.dict_params = nn.ParameterDict(
+            {
+                "foo": nn.Parameter((4,), "float32"),
+                "bar": nn.Parameter((4,), "float32"),
+            }
+        )
+
+
+def test_parameter_list_basic_behavior():
+    p0 = nn.Parameter((4,), "float32")
+    p1 = nn.Parameter((4,), "float32")
+    params = nn.ParameterList([p0])
+    params.append(p1)
+
+    assert len(params) == 2
+    assert params[0] is p0
+    assert list(params) == [p0, p1]
+
+    p2 = nn.Parameter((4,), "float32")
+    params[1] = p2
+    assert params[1] is p2
+
+    p3 = nn.Parameter((4,), "float32")
+    params.extend([p3])
+    assert list(params) == [p0, p2, p3]
+
+
+def test_parameter_dict_basic_behavior():
+    p0 = nn.Parameter((4,), "float32")
+    p1 = nn.Parameter((4,), "float32")
+    params = nn.ParameterDict({"foo": p0})
+    params["bar"] = p1
+
+    assert len(params) == 2
+    assert params["foo"] is p0
+    assert "bar" in params
+    assert list(params) == ["foo", "bar"]
+    assert list(params.keys()) == ["foo", "bar"]
+    assert list(params.values()) == [p0, p1]
+    assert list(params.items()) == [("foo", p0), ("bar", p1)]
+    assert params.get("foo") is p0
+
+    p2 = nn.Parameter((4,), "float32")
+    params.update({"baz": p2})
+    assert list(params.keys()) == ["foo", "bar", "baz"]
+    assert params.pop("baz") is p2
+    params.clear()
+    assert len(params) == 0
+
+
+def test_type_validation():
+    with pytest.raises(TypeError):
+        nn.ParameterList([object()])
+
+    with pytest.raises(TypeError):
+        nn.ParameterDict({"bad": object()})
+
+    with pytest.raises(TypeError):
+        nn.ParameterDict({1: nn.Parameter((4,), "float32")})
+
+    with pytest.raises(TypeError):
+        nn.ParameterList()[0] = object()
+
+
+def test_named_parameters_parameters_and_state_dict():
+    m = ParamContainerModule()
+
+    expected = [
+        "list_params.0",
+        "list_params.1",
+        "dict_params.foo",
+        "dict_params.bar",
+    ]
+
+    assert list(m.state_dict().keys()) == expected
+    assert [name for name, _ in m.named_parameters()] == expected
+    assert len(list(m.parameters())) == 4
+
+
+def test_nested_traversal_through_module_dict():
+    class Inner(nn.Module):
+        def __init__(self):
+            self.params = nn.ParameterList([nn.Parameter((4,), "float32")])
+
+    class Outer(nn.Module):
+        def __init__(self):
+            self.blocks = nn.ModuleDict({"inner": Inner()})
+
+    m = Outer()
+    assert list(m.state_dict().keys()) == ["blocks.inner.params.0"]
+
+
+def test_nested_traversal_through_module_list():
+    class Inner(nn.Module):
+        def __init__(self):
+            self.params = nn.ParameterList([nn.Parameter((4,), "float32")])
+
+    class Outer(nn.Module):
+        def __init__(self):
+            self.blocks = nn.ModuleList([Inner()])
+
+    m = Outer()
+    assert list(m.state_dict().keys()) == ["blocks.0.params.0"]
+
+
+def test_to_dtype():
+    m = ParamContainerModule()
+    m.to(dtype="float16")
+
+    assert m.list_params[0].dtype == "float16"
+    assert m.list_params[1].dtype == "float16"
+    assert m.dict_params["foo"].dtype == "float16"
+    assert m.dict_params["bar"].dtype == "float16"
+
+
+def test_load_state_dict():
+    m = ParamContainerModule()
+    p0 = nn.Parameter((4,), "float32")
+    p0.data = np.full((4,), 1.0, dtype="float32")
+    p1 = nn.Parameter((4,), "float32")
+    p1.data = np.full((4,), 2.0, dtype="float32")
+    p2 = nn.Parameter((4,), "float32")
+    p2.data = np.full((4,), 3.0, dtype="float32")
+    p3 = nn.Parameter((4,), "float32")
+    p3.data = np.full((4,), 4.0, dtype="float32")
+    state_dict = {
+        "list_params.0": p0,
+        "list_params.1": p1,
+        "dict_params.foo": p2,
+        "dict_params.bar": p3,
+    }
+
+    missing_keys, unexpected_keys = m.load_state_dict(state_dict)
+
+    assert missing_keys == []
+    assert unexpected_keys == []
+    tvm.testing.assert_allclose(m.list_params[0].data.numpy(), np.full((4,), 1.0, "float32"))
+    tvm.testing.assert_allclose(m.list_params[1].data.numpy(), np.full((4,), 2.0, "float32"))
+    tvm.testing.assert_allclose(m.dict_params["foo"].data.numpy(), np.full((4,), 3.0, "float32"))
+    tvm.testing.assert_allclose(m.dict_params["bar"].data.numpy(), np.full((4,), 4.0, "float32"))
+
+
+def test_export_tvm_parameter_names():
+    class M(nn.Module):
+        def __init__(self):
+            self.biases = nn.ParameterList(
+                [
+                    nn.Parameter((4,), "float32"),
+                    nn.Parameter((4,), "float32"),
+                ]
+            )
+            self.scales = nn.ParameterDict({"main": nn.Parameter((4,), "float32")})
+
+        def forward(self, x):
+            return x + self.biases[0] + self.biases[1] + self.scales["main"]
+
+    _, params = M().export_tvm(
+        spec={"forward": {"x": nn.spec.Tensor((4,), "float32")}},
+        debug=False,
+    )
+    assert [name for name, _ in params] == ["biases.0", "biases.1", "scales.main"]
+
+
+def test_mutator_parameter_container_names():
+    seen = []
+
+    class Recorder(nn.Mutator):
+        def visit_param(self, name: str, node: nn.Parameter) -> Any:
+            seen.append(name)
+            return node
+
+    m = ParamContainerModule()
+    Recorder().visit_module("", m)
+
+    assert seen == [
+        "list_params.0",
+        "list_params.1",
+        "dict_params.foo",
+        "dict_params.bar",
+    ]
+
+
+if __name__ == "__main__":
+    tvm.testing.main()