Repository: pytorch
Issue: torch.export.export fails to trace through a binary operator
Labels: triaged,oncall: pt2,export-triaged,oncall: export

### 🐛 Describe the bug

torch.export.export fails to trace through a binary operator.
Here is the failing code

```python
import torch

class MyTensor:
    def __init__(self, tensor):
        self.tensor = tensor

    def __mul__(self, rhs):
        return self.tensor * rhs

class MyModule(torch.nn.Module):
    def forward(self, t: torch.Tensor):
        my_tensor = MyTensor(torch.ones_like(t))
        return my_tensor * t # Fails to trace this

my_module = MyModule()

torch.export.export(my_module, args=(torch.tensor([2, 3]),))
```

This results in an error
```
Traceback (most recent call last):
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 645, in proxy_args_kwargs
    proxy_args = tuple(arg.as_proxy() for arg in args)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 645, in <genexpr>
    proxy_args = tuple(arg.as_proxy() for arg in args)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/variables/base.py", line 253, in as_proxy
    raise NotImplementedError(str(self))
NotImplementedError: UserDefinedObjectVariable(MyTensor)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/bpetkant/ws/scratchpad/scratchpad.py", line 17, in <module>
    torch.export.export(my_module, args=(torch.tensor([2, 3]),))
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/export/__init__.py", line 174, in export
    return _export(
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/export/_trace.py", line 945, in wrapper
    raise e
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/export/_trace.py", line 928, in wrapper
    ep = fn(*args, **kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/export/exported_program.py", line 89, in wrapper
    return fn(*args, **kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1455, in _export
    aten_export_artifact = export_func(
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/export/_trace.py", line 1060, in _strict_export
    gm_torch_level = _export_to_torch_ir(
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/export/_trace.py", line 512, in _export_to_torch_ir
    gm_torch_level, _ = torch._dynamo.export(
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1379, in inner
    result_traced = opt_f(*args, **kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 433, in _fn
    return fn(*args, **kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1116, in __call__
    return self._torchdynamo_orig_callable(
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 472, in __call__
    return _compile(
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_utils_internal.py", line 84, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
  File "/usr/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 817, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 636, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1185, in transform_code_object
    transformations(instructions, code_options)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 178, in _fn
    return fn(*args, **kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 582, in transform
    tracer.run()
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2451, in run
    super().run()
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 893, in run
    while self.step():
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 805, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 234, in impl
    self.push(fn_var.call_function(self, self.popn(nargs), {}))
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 962, in call_function
    return handler(tx, args, kwargs)
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/variables/builtin.py", line 907, in _handle_insert_op_in_graph
    *proxy_args_kwargs(args, kwargs),
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 652, in proxy_args_kwargs
    unimplemented(
  File "/home/bpetkant/ws/sharktank/.venv/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 220, in unimplemented
    raise Unsupported(msg) from from_exc
torch._dynamo.exc.Unsupported: call_function args: UserDefinedObjectVariable(MyTensor) TensorVariable() 

from user code:
   File "/home/bpetkant/ws/scratchpad/scratchpad.py", line 13, in forward
    return my_tensor * t

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
```

`MyTensor` is internal to the program and not part of the signature, so it should not require registration, but even registration with `torch.utils._pytree.register_pytree_node` does not help.

Note that the below code exports fine.
The difference is that instead of calling `__mul__` on `MyTensor`, it calls `my_mul`, which has the same function body.
```python
import torch

class MyTensor:
    def __init__(self, tensor):
        self.tensor = tensor

    def my_mul(self, rhs):
        return self.tensor * rhs

class MyModule(torch.nn.Module):
    def forward(self, t: torch.Tensor):
        my_tensor = MyTensor(torch.ones_like(t))
        return my_tensor.my_mul(t)

my_module = MyModule()

torch.export.export(my_module, args=(torch.tensor([2, 3]),))
```

### Versions

Collecting environment information...
PyTorch version: 2.4.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.12 (main, Mar 22 2024, 16:50:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-70-generic-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   52 bits physical, 57 bits virtual
Byte Order:                      Little Endian
CPU(s):                          192
On-line CPU(s) list:             0-191
Vendor ID:                       AuthenticAMD
Model name:                      AMD EPYC 9454 48-Core Processor
CPU family:                      25
Model:                           17
Thread(s) per core:              2
Core(s) per socket:              48
Socket(s):                       2
Stepping:                        1
Frequency boost:                 enabled
CPU max MHz:                     3810.7910
CPU min MHz:                     1500.0000
BogoMIPS:                        5492.22
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid overflow_recov succor smca fsrm flush_l1d
Virtualization:                  AMD-V
L1d cache:                       3 MiB (96 instances)
L1i cache:                       3 MiB (96 instances)
L2 cache:                        96 MiB (96 instances)
L3 cache:                        512 MiB (16 instances)
NUMA node(s):                    2
NUMA node0 CPU(s):               0-47,96-143
NUMA node1 CPU(s):               48-95,144-191
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] mypy==1.8.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.3
[pip3] onnx==1.15.0
[pip3] torch==2.4.1+cpu
[conda] Could not collect

cc @ezyang @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4