[AudioLM] Graph break: call_method UserDefinedObjectVariable(dict) get [TorchVariable(<class 'torch.Tensor'>), ConstantVariable(NoneType)]

[ezyang]

🐛 Describe the bug

Discovered while compiling lucidrains/audiolm-pytorch

[2024-03-06 11:44:09,519] [15/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG] Graph break: call_method UserDefinedObjectVariable(dict) get [TorchVariable(<class 'torch.Tensor'>), ConstantVariable(NoneType)] {} from user code at:
[2024-03-06 11:44:09,519] [15/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/home/ezyang/local/miniconda3-test/envs/audiolm/lib/python3.10/site-packages/einops-0.7.0-py3.10.egg/einops/packing.py", line 141, in resume_in_unpack
[2024-03-06 11:44:09,519] [15/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     backend = get_backend(tensor)
[2024-03-06 11:44:09,519] [15/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]   File "/home/ezyang/local/miniconda3-test/envs/audiolm/lib/python3.10/site-packages/einops-0.7.0-py3.10.egg/einops/_backends.py", line 28, in get_backend
[2024-03-06 11:44:09,519] [15/0] torch._dynamo.symbolic_convert.__graph_breaks: [DEBUG]     _result = _type2backend.get(_type, None)

This is a bit concerning; if it's actually a dict, we should have wrapped it in UserDefinedObjectVariable. This lookup is happening in einops which had some problems

Full repro code: gist.github.com/ezyang/64c24c9fc5529f3afed4ee4266f6adc5

Versions

main

cc @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng

[ezyang]

Looking over the einops code, it looks like this is a normal dict:

_loaded_backends: dict = {}
_type2backend: dict = {}
_debug_importing = False


def get_backend(tensor) -> "AbstractBackend":
    """
    Takes a correct backend (e.g. numpy backend if tensor is numpy.ndarray) for a tensor.
    If needed, imports package and creates backend
    """
    _type = type(tensor)
    _result = _type2backend.get(_type, None)
    if _result is not None:
        return _result

    for framework_name, backend in list(_loaded_backends.items()):
        if backend.is_appropriate_type(tensor):
            _type2backend[_type] = backend
            return backend

[Skylion007]

I am having some similar problems with nested dicts right now. Except it's even worse as it's leading to an uncaught assestion in the constants.py (passing a Dict that should have been a ConstDict).

[yanboliang]

@Skylion007 Do you have a small repro for your issue? I think we should find out the scenario that a regular is not wrapped as ConstDict, then it's not hard to fix it imo.

[ezyang]

If we just add allocation stack traces to all variable trackers this will be easy to nail. If someone gives me a patch I can run it on AudioLM

[lezcano]

Dictionaries in dynamo have some limitations. The good part of the story is that these limitations are documented and it's super easy to extend our implementation to support more types!

pytorch/torch/_dynamo/variables/dicts.py

Lines 27 to 29 in 06d2392

	# [Adding a new supported class within the keys of ConstDictVarialble]
	# - Add its tracker type to is_hashable
	# - (perhaps) Define how it is compared in _HashableTracker._eq_impl

At the moment we are graph-breaking when hitting elements of dicts that are not hashable internally:

pytorch/torch/_dynamo/variables/dicts.py

Lines 74 to 75 in 06d2392

	if not is_hashable(vt):
	unimplemented(f"Dict key of type {type(vt)}. Key: {vt}")

We could throw a hard error saying "this dict key is not supported, please file a ticket with your error to ask for support for this type of keys".