torch/_refs/__init__.py is not autoreloadable in bento / jupyter notebook

[ezyang]

🐛 Describe the bug

When we try to reload this file, it fails with this exception:

 File "/mnt/xarfuse/uid-128580/66551cbc-seed-nspid4026531836_cgpid58278121-ns-4026531840/parsh/autoreload/autoreload.py", line 627, in _reload_modules
    self.superreload(m)
  File "/mnt/xarfuse/uid-128580/66551cbc-seed-nspid4026531836_cgpid58278121-ns-4026531840/parsh/autoreload/autoreload.py", line 836, in superreload
    self._superreload_helper(module)
  File "/mnt/xarfuse/uid-128580/66551cbc-seed-nspid4026531836_cgpid58278121-ns-4026531840/parsh/autoreload/autoreload.py", line 773, in _superreload_helper
    module = reload(module)
  File "/mnt/xarfuse/uid-128580/66551cbc-seed-nspid4026531836_cgpid58278121-ns-4026531840/runtime/lib/python3.10/importlib/__init__.py", line 203, in reload
    _bootstrap._exec(spec, module)
  File "<frozen importlib._bootstrap>", line 778, in _exec
  File "/mnt/xarfuse/uid-128580/66551cbc-seed-nspid4026531836_cgpid58278121-ns-4026531840/parsh/autoreload/measurements.py", line 46, in patched_exec_module
    orig_exec_module(self_, module)
  File "<frozen importlib._bootstrap_external>", line 896, in exec_module
  File "<frozen importlib._bootstrap>", line 400, in _call_with_frames_removed
  File "/data/users/ezyang/fbsource/fbcode/caffe2/torch/_refs/__init__.py", line 526, in <module>
    def abs(a):
  File "/data/users/ezyang/fbsource/fbcode/caffe2/torch/_refs/__init__.py", line 475, in inner
    aten_op = utils.get_aten_op(prim, prim.__name__)
AssertionError
]

The proximal cause seems to be that we are trying to reregister all of the operators, but the old registrations didn't get cleared on unload.

cc @anjali411 @zou3519

Versions

main