Method test_prune_tree

axlearn/common/utils_test.py:947–972

(self)

            check_jax_type(kwargs={"key": "1"})
        with self.assertRaisesRegex(ValueError, "^Argument key has leaf with non-JAX type"):
            check_jax_type(pretty_named_args={"key": "1"})

    def test_prune_tree(self):
        in_tree = {
            "a": {
                "b": {"d": "test"},
                "c": {
                    "b": None,
                    "e": VDict({"ee": 123}),
                },
            },
            "f": 345,
        }
        # Prune by path.
        result = prune_tree(in_tree, lambda k, _: "b" in k)
        self.assertEqual({"a": {"c": {"e": VDict({"ee": 123})}}, "f": 345}, result)
        # VDict should be preserved.
        self.assertIsInstance(result["a"]["c"]["e"], VDict)
        # Prune by path with prefix/separator.
        self.assertEqual(
            {"a": {"c": {"b": None, "e": {"ee": 123}}}, "f": 345},
            prune_tree(in_tree, lambda k, _: k == "prefix:a:b", prefix="prefix", separator=":"),
        )
        # Prune by value.
        self.assertEqual(
            {"a": {"b": {"d": "test"}, "c": {"b": None, "e": VDict()}}},
            prune_tree(in_tree, lambda _, v: isinstance(v, int)),
        )