import functorch fails (cannot import name 'set_autograd_function_allowed' from 'torch._C._functorch')

[kvoronin-intel]

Hi!

I'm new to IREE and tried to follow instructions at https://iree-org.github.io/iree/getting-started/pytorch/ to run a simple example provided there, https://github.com/iree-org/iree-torch/blob/main/examples/regression.py.

However, and even when I run regression.py, or even smth as simple as import functorch I am getting an error:

>>> import functorch
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/nfs_home/kvoronin/anaconda/envs/iree_try_venv/lib/python3.10/site-packages/functorch/__init__.py", line 16, in <module>
    from ._src.vmap import vmap
  File "/nfs_home/kvoronin/anaconda/envs/iree_try_venv/lib/python3.10/site-packages/functorch/_src/vmap/__init__.py", line 4, in <module>
    from torch._functorch.vmap import (
  File "/nfs_home/kvoronin/anaconda/envs/iree_try_venv/lib/python3.10/site-packages/torch/_functorch/vmap.py", line 24, in <module>
    from torch._functorch.utils import exposed_in
  File "/nfs_home/kvoronin/anaconda/envs/iree_try_venv/lib/python3.10/site-packages/torch/_functorch/utils.py", line 3, in <module>
    from torch._C._functorch import (
ImportError: cannot import name 'set_autograd_function_allowed' from 'torch._C._functorch' (unknown location)

Python 3.10.8 torch 2.0.0.dev20230106+cpu print(torch.version) 1.13.1

I guess it is some version mismatch, but the set_autograd_function_allowed is badly google-able. Any advice would be welcome!

[kvoronin-intel]

If on top of the described setup, I do pip install functorch I get, while updating,

ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
torch-mlir 20230108.712 requires torch==2.0.0.dev20230106, but you have torch 1.13.1 which is incompatible.  

and then, when running the regression.py,

(iree_try_venv) bash-4.2$ python -u example_regression.py 
Traceback (most recent call last):
  File "/data/nfs_home/kvoronin/work/tensor_compiler/iree/example/example_regression.py", line 7, in <module>
    import torch_mlir
  File "/nfs_home/kvoronin/anaconda/envs/iree_try_venv/lib/python3.10/site-packages/torch_mlir/__init__.py", line 12, in <module>
    from torch._functorch.compile_utils import strip_overloads
ModuleNotFoundError: No module named 'torch._functorch'

so it looks to an inexperienced person like me that something is off between torch_mlir, torch and functorch versions.

What is the working combination then?

[dellis23]

Hey @kvoronin-intel, your intuition is likely correct. We frequently have issues with version mismatches. We're in the process of discussing how we can get this aligned across projects, but that hasn't happened yet. My personal fix for this has been to go back to a known good version (perhaps via the tests on this page) and install the nightlies for the various packages from that date (i.e. torch-mlir, iree, and the pytorch version they depend on).

I'm sorry this is difficult right now, but hopefully we'll have a fix in the future. This is definitely not the experience we want for users.