Motivation

Google Colab recently switched from Python 3.11 to Python 3.12, which introduced several dependency issues, particularly affecting our TensorFlow backend.

The main issue was that we had pinned TensorFlow to version 2.12.0 (tensorflow>=2.0,<=2.12 for non-Windows, <=2.10 for Windows). These old versions are no longer available for Python 3.12, where the earliest installable TensorFlow release is 2.16.0.

However, since TensorFlow 2.16.0, Keras 3 is used by default, which is incompatible with our current TensorFlow code (e.g., BatchNormalization is not available with Keras 3 errors due to tf.compat.v1). This incompatibility was the reason for pinning TensorFlow in the first place.

Changes Introduced

To resolve these issues, this PR introduces the following changes:

1. TensorFlow & Keras compatibility (link to TF documentation)

• Unpinned TensorFlow, allowing installation of recent versions.
• Forced the use of Keras 2 (tf-keras) by installing tf-keras and setting the environment variable TF_USE_LEGACY_KERAS=1.
• Introduced a new module _tf_legacy.py to ensure these environment variables are declared when using the TensorFlow backend.

2. Additional environment variables

• Added WRAPT_DISABLE_EXTENSIONS=1 to resolve the error:
  TypeError: this __dict__ descriptor does not support '_DictWrapper' objects
  (related to wrapt, a TensorFlow dependency, the solution has been suggested here).

3. Compatibility shim for tf-keras namespace changes

• Newer versions of tf-keras reorganized legacy layers under a src/ namespace.
• Added a shim in _tf_legacy.py to handle both layouts:

try:
    import tf_keras.src.legacy_tf_layers as legacy_tf_layers
    sys.modules["tf_keras.legacy_tf_layers"] = legacy_tf_layers
except ImportError:
    # Older tf-keras didn’t use src/, so nothing to do
    pass

4. Relaxed other package requirements

• Relaxed pinned versions for tables and PySide6, since the exact pinned versions were unavailable on Python 3.12.

5. CI/CD improvements

• Extended the test matrix to include Python 3.11 and Python 3.12 on Linux, MacOS, and Windows.
• On MacOS + Python 3.12, functional tests were hanging due to a known PyTorch DataLoader bug when num_workers > 0. Set num_workers=0 explicitly in these cases.
• Attempted to include Python 3.13, but encountered multiple unrelated issues across OSs. Deferred this for later (may become necessary once Colab updates to 3.13).

6. Colab notebooks update

• Simplified installation steps in TensorFlow-based Colab notebooks. GPU access now works without the workaround previously required.

Testing

• Verified unit and functional tests pass across the updated CI/CD matrix.
• Ran the 3 TensorFlow-based notebooks manually — all executed successfully.
• ⚠️ MacOS with tensorflow-metal: only tested via CI/CD (no tensorflow-metal), not on a real machine. This should be validated before merging.

Notes / Future Work

• Monitor TensorFlow and Colab’s transition to Python 3.13 and plan for compatibility.
• Ensure MacOS + tensorflow-metal users validate installation and training pipelines.