Errors with incompatible versions of torch_xla and ai-edge-torch

[ccleavinger]

Description of the bug:

Got an ImportError from ai-edge-torch when I imported ai-edge torch. Code:

# the variable model is a predefined pytorch model that I imported from a state dict.
# torch and other imports are declared earlier
import ai_edge_torch as aie

sample_inputs = (torch.randn(1,3,448,448))

edge_model = aie.convert(model, sample_inputs)
edge_model.export('lite_model.tflite')

Error:

ImportError                               Traceback (most recent call last)
[<ipython-input-7-98c1a8b76fa9>](https://localhost:8080/#) in <cell line: 1>()
----> 1 import ai_edge_torch as aie
      2 
      3 sample_inputs = (torch.randn(1,3,448,448))
      4 
      5 edge_model = aie.convert(model, sample_inputs)

4 frames
[/usr/local/lib/python3.10/dist-packages/ai_edge_torch/__init__.py](https://localhost:8080/#) in <module>
     14 # ==============================================================================
     15 
---> 16 from .convert.converter import convert
     17 from .convert.converter import signature
     18 from .model import Model

[/usr/local/lib/python3.10/dist-packages/ai_edge_torch/convert/converter.py](https://localhost:8080/#) in <module>
     20 import torch
     21 
---> 22 from ai_edge_torch import model
     23 from ai_edge_torch.convert import conversion
     24 from ai_edge_torch.convert import conversion_utils as cutils

[/usr/local/lib/python3.10/dist-packages/ai_edge_torch/model.py](https://localhost:8080/#) in <module>
     26 import tensorflow as tf
     27 
---> 28 from ai_edge_torch.convert import conversion_utils as cutils
     29 
     30 

[/usr/local/lib/python3.10/dist-packages/ai_edge_torch/convert/conversion_utils.py](https://localhost:8080/#) in <module>
     23 
     24 import torch
---> 25 from torch_xla import stablehlo
     26 
     27 from ai_edge_torch.quantize import quant_config as qcfg

[/usr/local/lib/python3.10/dist-packages/torch_xla/__init__.py](https://localhost:8080/#) in <module>
      5 
      6 import torch
----> 7 import _XLAC
      8 from ._internal import tpu
      9 

ImportError: /usr/local/lib/python3.10/dist-packages/_XLAC.cpython-310-x86_64-linux-gnu.so: undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKSs

Actual vs expected behavior:

I expected a tf-lite model to be exported.

Any other information you'd like to share?

My version of pytorch is 2.4.0 (as this library requires) and torch_xla is 2.3.* I believe this is where the issue stems from. I simply don't know where a version of torch_xla that is compatible with the 2.4 version of pytorch can be installed. Sorry if this is a dumb issue I just desperately need to fix this. Any help would be appreciated!

[chunnienc]

Hi @ccleavinger , please follow our setup guide to install the supported nightly dependencies. Here are the commands:

pip uninstall torch torchvision torch_xla

pip install -r https://raw.githubusercontent.com/google-ai-edge/ai-edge-torch/main/requirements.txt
pip install ai-edge-torch-nightly

and make sure your installed dependencies match the versions specified here https://github.com/google-ai-edge/ai-edge-torch/blob/main/requirements.txt . Please let me know if you have any issues with the setup, thanks!

[ccleavinger]

Thank you I just had some library conflicts! I should've used the requirements listed on the github I just thought they would've been handled on the PyPi version. Thanks so much for the help you're a lifesaver!