[JIT] nested dictionaries are not traced correctly. #75012
Comments
facebook-github-bot
added
the
oncall: jit
|
Hi. I have the same problem. have you found any solution? |
|
@nimaous Currently I just flatten and unflatten the dictionary as needed from typing import Any, Callable, Sequence
def flatten_dict(
d: dict[str, Any],
/,
*,
recursive: bool = True,
join_fn: Callable[[Sequence[str]], str] = ".".join,
) -> dict[str, Any]:
r"""Flatten dictionaries recursively."""
result: dict[str, Any] = {}
for key, item in d.items():
if isinstance(item, dict) and recursive:
subdict = flatten_dict(item, recursive=True, join_fn=join_fn)
for subkey, subitem in subdict.items():
result[join_fn((key, subkey))] = subitem
else:
result[key] = item
return result
def unflatten_dict(
d: dict[str, Any],
/,
*,
recursive: bool = True,
split_fn: Callable[[str], Sequence[str]] = lambda s: s.split(".", maxsplit=1),
) -> dict[str, Any]:
r"""Unflatten dictionaries recursively."""
result = {}
for key, item in d.items():
split = split_fn(key)
result.setdefault(split[0], {})
if len(split) > 1 and recursive:
assert len(split) == 2
subdict = unflatten_dict(
{split[1]: item}, recursive=recursive, split_fn=split_fn
)
result[split[0]] |= subdict
else:
result[split[0]] = item
return result
a = {
"a": True,
"b": 42,
"c": "foo",
"foo": {"a": 1, "b": 2, "c": 3},
"bar": {"a": {"a": 1}, "b": {"a": 1}, "c": {"a": 1}},
"baz": {"a": True, "b": 2}, # <-- does not work!
}
print(flat := flatten_dict(a))
print(unflat := unflatten_dict(a))
assert unflat == a |
|
Same problem, will pytorch team support this? |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
馃悰 Describe the bug
Second and third level dictionaries are only accepted if they are all the same type.
Raises
Which is weird because top level dictionaries with variable types get traced just fine. Commenting
bazand it compiles fine as well as then all nested dictionaries are of the same type.Versions
The text was updated successfully, but these errors were encountered: