Method test_tree_paths

axlearn/common/utils_test.py:203–284

(self)

        self.assertEqual(cfg.layer[1].input_dim, updated_input_dim)


class TreeUtilsTest(TestCase):
    def test_tree_paths(self):
        tree = {"a": 1, "b": [2, {"c": 3}]}
        self.assertEqual({"a": "a", "b": ["b/0", {"c": "b/1/c"}]}, tree_paths(tree))

        # Tuple.
        self.assertEqual(("0", ("1/0", "1/1"), "2"), tree_paths(("a", ("b", "c"), "d")))

        # NamedTuple.
        self.assertEqual(
            Combo(head="head", tail=Combo(head="tail/head", tail="tail/tail")),
            tree_paths(Combo(head=1, tail=Combo(head=2, tail=3))),
        )

        # flax_struct.PyTreeNode.
        self.assertEqual(
            WeightedSummary(mean="mean", weight="weight"),
            tree_paths(WeightedSummary(mean=2, weight=3)),
        )

        # Nested flax_struct.PyTreeNode.
        self.assertEqual(
            StructContainer(WeightedSummary(mean="contents/mean", weight="contents/weight")),
            tree_paths(StructContainer(WeightedSummary(mean=2, weight=3))),
        )

        # str-Enum key.
        class MyEnum(str, enum.Enum):
            RED = "red"

        self.assertEqual({MyEnum.RED: "MyEnum.RED"}, tree_paths({MyEnum.RED: 3}))

        # With is_leaf set.
        self.assertEqual(
            ["0", {"a": "1/a", "b": "1/b"}],
            tree_paths(
                [Combo(head=1, tail=2), {"a": Combo(head=3, tail=4), "b": 5}],
                is_leaf=lambda x: isinstance(x, Combo),
            ),
        )

        class DataclassCombo(flax_struct.PyTreeNode):
            scalar: int
            dataclass_combo: Any
            none: type[None]
            nested_tensor: NestedTensor

        # Nested custom pytree.
        self.assertEqual(
            DataclassCombo(  # pytype: disable=wrong-arg-types
                scalar="scalar",
                dataclass_combo=DataclassCombo(  # pytype: disable=wrong-arg-types
                    scalar="dataclass_combo/scalar",
                    dataclass_combo=Combo(
                        head="dataclass_combo/dataclass_combo/head",
                        tail="dataclass_combo/dataclass_combo/tail",
                    ),