AttributeError: module 'tensorflow' has no attribute 'saved_model' when importing gemma3 from ai_edge_torch in Google Colab

Question

I'm trying to use Google's ai-edge-torch library to work with the Gemma 3 model in a Google Colab environment. However, I encounter the following error during import:

import torch
from ai_edge_torch.generative.examples.gemma3 import gemma3

Error traceback:

/content/ai-edge-torch/ai_edge_torch/model.py in <module>
 31 from ai_edge_litert import interpreter as tfl_interpreter # pylint: disable=g-direct-tensorflow-import
 32 
---> 33 DEFAULT_SIGNATURE_NAME = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
 34 
AttributeError: module 'tensorflow' has no attribute 'saved_model'

This suggests that the tensorflow module being used does not expose the saved_model attribute, which is unexpected since tf.saved_model has been available since TensorFlow 2.0.

I'm running this in a Google Colab notebook (as of December 2025). I’ve tried restarting the runtime and installing the latest versions of the relevant packages, but the issue persists.

Environment info:

• Google Colab

• TensorFlow version: (output of tf.__version__ — possibly outdated or conflicting)

• ai-edge-torch version: (e.g., installed via pip install ai-edge-torch)

• Sagar: This appears to be a version conflict or an incomplete TensorFlow installation. Upgrade or reinstall TensorFlow, then restart the runtime to resolve the missing saved_model error.