Dynamo does not support methods defined on NamedTuples

[suo]

from typing import NamedTuple

class MyTuple(NamedTuple):
    first: torch.Tensor
    second: torch.Tensor

    def add(self):
        return first + second

def my_fn(x):
    return x.add()

torch._dynamo.export(my_fn, MyTuple(first=torch.tensor(1), second=torch.tensor(2)))

raises torch._dynamo.exc.Unsupported: NamedTupleVariable.add

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh @mlazos @voznesenskym @yanboliang @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @chunyuan-w @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

[suo]

seems a relatively simple fix, we support methods on user defined classes after all

[ShisuiUzumaki]

@suo would you be accepting a PR for this? If so, and it is not too difficult to implement, I am interested, though would require instructions from your end

[ShisuiUzumaki]

@suo @zou3519 Could you give some initial instructions on how to tackle this issue

[ezyang]

Need to trigger inlining translation on the method call, symmetric with how UserDefinedObject can inline. Also cc @wconstab who has worked on inlining

[wconstab]

I could help guide you @ShisuiUzumaki, unless you are already connecting with others on the thread. You can ping me on pytorch slack for more 1:1, and here are some starting pointers.

1. the NamedTupleVariable class is dynamo's way of modeling behaviors that are supported on this type of object if it's encountered during tracing. It needs to implementcall_method.
2. NamedTupleVariable already extends from TupleVariable which implements`call_method'. So it may be that you just want to extend TupleVariable's call_method to be more general.
3. TupleVariable extends from BaseListVariable and that also implements call_method in a slightly more general way than TupleVariable, and then delegates to the ultimate base class 'VariableTracker' which handles a couple more cases and finally raises Unimplemented.

The trivial part of this PR would be just adding general support to handle 'call_method' at the NamedTupleVariable layer.
For this part you can take inspiration from how 'UnspecializedNNModuleVariable.call_method' calls inline_user_function_return after constructing a VariableTracker for the Function object.

One thing I'm not sure about is, how much of the generality should be allowed only for NamedTuple vs Tuple vs List? Is a user allowed to define a custom List with methods, or not? cc @suo

[suo]

I think NamedTuple is the only one we really need to support atm. Theoretically you could inherit from List and define some methods, but no one really does that. In contrast, it's fairly common to hang some methods off your NamedTuple.

[ShisuiUzumaki]

I could help guide you @ShisuiUzumaki, unless you are already connecting with others on the thread. You can ping me on pytorch slack for more 1:1, and here are some starting pointers.

1. the NamedTupleVariable class is dynamo's way of modeling behaviors that are supported on this type of object if it's encountered during tracing. It needs to implementcall_method.
2. NamedTupleVariable already extends from TupleVariable which implements`call_method'. So it may be that you just want to extend TupleVariable's call_method to be more general.
3. TupleVariable extends from BaseListVariable and that also implements call_method in a slightly more general way than TupleVariable, and then delegates to the ultimate base class 'VariableTracker' which handles a couple more cases and finally raises Unimplemented.

The trivial part of this PR would be just adding general support to handle 'call_method' at the NamedTupleVariable layer. For this part you can take inspiration from how 'UnspecializedNNModuleVariable.call_method' calls inline_user_function_return after constructing a VariableTracker for the Function object.

One thing I'm not sure about is, how much of the generality should be allowed only for NamedTuple vs Tuple vs List? Is a user allowed to define a custom List with methods, or not? cc @suo

@wconstab I went through the contents of call_method. I got the just of what it is doing but would appreciate it if you could give an in-depth explanation, or more like the intuition behind it

[ShisuiUzumaki]

@wconstab do you think something like this might work

  def call_method(
      self,
      tx,
      name,
      args: "List[VariableTracker]",
      kwargs: "Dict[str, VariableTracker]",
  ) -> "VariableTracker":
      options = VariableTracker.propagate(self, args, kwargs.values())
      if (
          name in ("__add__", "__iadd__")
          and len(args) == 1
          and isinstance(args[0], NamedTupleVariable)
      ):
          assert not kwargs
          return NamedTupleVariable(
              [x + y for x, y in zip(self.items, args[0].items)],
              self.tuple_cls,
              **options,
          )
      elif (
          name in ("__add__", "__iadd__")
          and len(args) == 1
          and isinstance(args[0], variables.ConstantVariable)
      ):
          assert not kwargs
          return NamedTupleVariable(
              [x + args[0] for x in self.items],
              self.tuple_cls,
              **options,
          )
      return super(NamedTupleVariable, self).call_method(tx, name, args, kwargs)

[wconstab]

@wconstab do you think something like this might work

This seems like an attempt to define a new 'add' behavior for dynamo that is specific to a NamedTupleVariable + a constant.

I understood this PR as a more general need to support any type of method a user might have written on a class that derives from NamedTupleVariable.

The way I would envision it looking is more like

elif hasattr(self.value, name):   # does our object actually have this method that we're trying to call?
   method = inspect.getattr_static(type(self.value), name)   # get the actual method off our actual obj
   
   # now, you have to consider multiple cases where 'method' is a staticmethod, a function, etc. and implement these differently.  (See my hint above for how a staticmethod can be handled; slight variation would be needed for a bound function).
   
    if isinstance(method, staticmethod):
        # the source helps dynamo keep track of what real py obj we are working with
        source = AttrSource(
            AttrSource(AttrSource(self.source, "__class__"), name), "__func__"
        )
        
        # the inline call tells dynamo to start tracing more code
        return tx.inline_user_function_return(
            # the UserFunctionVariable is an object that models the underlying py function and lets dynamo track interactions with it.  inline_user_function will be tracing through the code of `method.__func__` in this case
            variables.UserFunctionVariable(
                method.__func__, source=source, **options
            ),
            args,
            kwargs,
        )

[ShisuiUzumaki]

@wconstab Are you referring to the elif part of the code I mentioned earlier? Do you think the If part is ok?

I understand somewhat from the hint you gave earlier, but I need more context regarding the primary purpose of call_method.

Also, in the code you provided above, you check for the self.method. value, is it checking in user-defined class which is inheriting from NamedTupleVariable