diff --git a/modelopt/torch/utils/_pytree.py b/modelopt/torch/utils/_pytree.py index 586c92b893e..7ff4078f4f7 100644 --- a/modelopt/torch/utils/_pytree.py +++ b/modelopt/torch/utils/_pytree.py @@ -114,6 +114,10 @@ def flatten_tree(pytree: Any, prefix: str = "") -> tuple[list[Any], TreeSpec]: Returns: A tuple (values, pytree) where values is a list of values flattened from the provided pytree, and tree_spec is the pytree spec describing the structure of the pytree. + + Raises: + ValueError: If two leaves map to the same flattened key (nested keys + are joined with ".", so ``{"a": {"b": 1}, "a.b": 2}`` collides). """ def collect_spec(pytree, prefix): @@ -128,6 +132,16 @@ def collect_spec(pytree, prefix): yield prefix, pytree # retrieve flattened values and names. Then initialize tree_spec with the flattened names. - flattened = dict(collect_spec(pytree, prefix)) - - return list(flattened.values()), TreeSpec(pytree, list(flattened.keys())) + seen_names: set[str] = set() + names: list[str] = [] + values: list[Any] = [] + for name, value in collect_spec(pytree, prefix): + if name in seen_names: + raise ValueError( + f"Cannot flatten pytree: multiple leaves map to the flattened key {name!r}!" + ) + seen_names.add(name) + names.append(name) + values.append(value) + + return values, TreeSpec(pytree, names) diff --git a/tests/unit/torch/utils/test_pytree.py b/tests/unit/torch/utils/test_pytree.py index 2dd95265528..4b6e7e9dee4 100644 --- a/tests/unit/torch/utils/test_pytree.py +++ b/tests/unit/torch/utils/test_pytree.py @@ -72,3 +72,31 @@ def test_flatten_tree(data, expected_keys): # try re-building to see if we get same structure tree_rebuilt = unflatten_tree(copy.deepcopy(values), tree_spec) assert data == tree_rebuilt + + +@pytest.mark.parametrize( + "data", + [ + {"a": {"b": 1}, "a.b": 2}, + {"a": [1], "a.0": 2}, + {"a.b": 1, "a": {"b": 2}}, + ], +) +def test_flatten_tree_key_collision(data): + """Colliding flattened keys must raise instead of silently dropping values.""" + with pytest.raises(ValueError, match=r"'a\.(b|0)'"): + flatten_tree(data) + + +@pytest.mark.parametrize( + ("data", "expected_keys"), + [ + ({"a": {"b": 1}, "ab": 2}, ["a.b", "ab"]), + ({"a": {"b": 1}, "a_b": 2}, ["a.b", "a_b"]), + ], +) +def test_flatten_tree_near_collision(data, expected_keys): + """Distinct flattened keys that merely look similar must still round-trip.""" + values, tree_spec = flatten_tree(data) + assert expected_keys == tree_spec.names + assert data == unflatten_tree(copy.deepcopy(values), tree_spec)