Defining custom torch function causes "maximum recursion depth" in autoreload

[willtalmadge]

Python 3.12.2
ipython 8.22.2
torch 2.2.1

I have a reproduction that is 100% reliable below.

Create a script with the following contents named test_module.py:

import torch

class CustomTorchFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input1, input2, weights):
        return 1

    @staticmethod
    def backward(ctx, grad_output):
        return 1


def main():
    var = 'something to edit 0'
    print("called main " + var)

Start ipython in the directory with the module

Execute:

%load_ext autoreload
%autoreload 2
import test_module
test_module.main()

It should print out

called main something to edit 0

So far so good. Now, edit test_module.py and increment 0 to 1. In ipython execute:

test_module.main()

The following error is printed out

[autoreload of test_module failed: Traceback (most recent call last):
  File "/Users/will/miniforge3/envs/ml0524/lib/python3.12/site-packages/IPython/extensions/autoreload.py", line 276, in check
    superreload(m, reload, self.old_objects)
  File "/Users/will/miniforge3/envs/ml0524/lib/python3.12/site-packages/IPython/extensions/autoreload.py", line 500, in superreload
    update_generic(old_obj, new_obj)
  File "/Users/will/miniforge3/envs/ml0524/lib/python3.12/site-packages/IPython/extensions/autoreload.py", line 397, in update_generic
    update(a, b)
  File "/Users/will/miniforge3/envs/ml0524/lib/python3.12/site-packages/IPython/extensions/autoreload.py", line 349, in update_class
    if update_generic(old_obj, new_obj):
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/will/miniforge3/envs/ml0524/lib/python3.12/site-packages/IPython/extensions/autoreload.py", line 397, in update_generic
    update(a, b)
  File "/Users/will/miniforge3/envs/ml0524/lib/python3.12/site-packages/IPython/extensions/autoreload.py", line 349, in update_class
    if update_generic(old_obj, new_obj):
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/will/miniforge3/envs/ml0524/lib/python3.12/site-packages/IPython/extensions/autoreload.py", line 397, in update_generic
    update(a, b)
  File "/Users/will/miniforge3/envs/ml0524/lib/python3.12/site-packages/IPython/extensions/autoreload.py", line 349, in update_class
    if update_generic(old_obj, new_obj):
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/will/miniforge3/envs/ml0524/lib/python3.12/site-packages/IPython/extensions/autoreload.py", line 397, in update_generic
    update(a, b)
  File "/Users/will/miniforge3/envs/ml0524/lib/python3.12/site-packages/IPython/extensions/autoreload.py", line 349, in update_class
    if update_generic(old_obj, new_obj):
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RecursionError: maximum recursion depth exceeded
]
called main something to edit 1

It seems like the module was successfully reloaded, but the error says it failed. So I'm not sure if I'm in an undefined state or not.

[willtalmadge]

For folks looking for a partial workaround, move the custom function into its own module. So the main code you are editing doesn't trigger this code to be reloaded. That seems to resolve this for me. Of course, if you need to edit the custom torch function, you may need to restart ipython or manually reload the module.

[Carreau]

thanks for the report, I'll see what I can do.