Dynamo: how to deal with multiple inheritance (nn.Module/MutableMapping)?

[vmoens]

🐛 Describe the bug

TensorDict is a MutableMapping object, and is treated as such by torch.compile:

import torch
from tensordict import TensorDict

td = TensorDict(a=1, b=2, c=True)

@torch.compile(fullgraph=True)
def add1(td):
    return TensorDict(**td)+1

add1(td)

We also have a TensorDictParams primitive that acts a bit like ParameterList: it is a TensorDict but also an nn.Module. That's useful when you want to set a TensorDict in an nn.Module have have the leaf tensors included in the state_dict, or dispatch ops like module.to(...) to the tensors it contains. However, _dynamo looks at it like an nn.Module and not a MutableMapping

import torch
from tensordict import TensorDictParams, TensorDict

td = TensorDictParams(TensorDict(a=1, b=2, c=True))

@torch.compile(fullgraph=True)
def add1(td):
    return TensorDict(**td)+1

add1(td)

breaks with

  File "/Users/vmoens/venv/rl/lib/python3.10/site-packages/torch/_dynamo/variables/dicts.py", line 357, in call_method
    dict_vt = BuiltinVariable.call_custom_dict(tx, dict, args[0])
  File "/Users/vmoens/venv/rl/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 1432, in call_custom_dict
    unimplemented(f"{user_cls.__name__}(): {args} {kwargs}")
  File "/Users/vmoens/venv/rl/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 313, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: dict(): (UnspecializedNNModuleVariable(TensorDictParams),) {}

My understanding is that call_custom_dict looks at the arg an in one case it's a variables.MutableMappingVariable which is fine but in the other it's a UnspecializedNNModuleVariable which isn't a mutable mapping.

So I guess my question is (other than how can we fix this) how does dynamo look at multiple inheritance? Shouldn't there be a way to tell "look, this isn't a bird or a fish but a fish that can fly"?

(note that in this specific case, smth(**obj) will call obj.keys() followed by obj.__getitem__ which are ops that compile is happy about - maybe that's what call_custom_dict should be doing?)

Here is a MRE:

import torch
from torch import nn
import collections

# class MyWeirdDict(collections.abc.MutableMapping):  # Works
class MyWeirdDict(collections.abc.MutableMapping, nn.Module):  # breaks
    def __init__(self, **kwargs):
        super().__init__()
        self._items = kwargs
    def keys(self):
        return self._items.keys()
    def __getitem__(self, item):
        return self._items[item]
    def __setitem__(self, key, value):
        self._items[key] = value
    def __delitem__(self, item):
        del self._items[item]
    def __len__(self):
        return len(self._items)
    def __iter__(self):
        yield from self._items
    def __hash__(self):
        return hash(id(self))
    def items(self):
        for k, v in self._items.items():
            yield (k, v)

@torch.compile(fullgraph=True)
def to_weird_dict(td):
    return MyWeirdDict(**td)

d = MyWeirdDict(a=1, b=2, c=3)
to_weird_dict(d)

Error logs

See above

Versions

nightlies

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames

[StrongerXi]

So I guess my question is (other than how can we fix this) how does dynamo look at multiple inheritance? Shouldn't there be a way to tell "look, this isn't a bird or a fish but a fish that can fly"?

@vmoens It doesn't seem to be well supported in general, but I think we need more datapoints to evaluate, so I created #142414 to track that.