Support Keras 3

[mhucka]

TensorFlow Quantum currently supports only Keras 2. The broader TensorFlow ecosystem has transitioned from Keras version 2 to Keras version 3. It would be good to update TFQ's Keras support to work with Keras 3.

[rosspeili]

Hey @mhucka happy to give this a shot. Where would you point at as a starting point? I will read keras 3 docs and compare to 2, review how it's curretnly used across the repo, and suggest a short pr if assigned. Thanks.

[mhucka]

@rosspeili Thanks for offering. I've assigned it to you.

I regret I'm not familiar enough with the frameworks to have a good recommendation, unfortunately. My naive approach would probably be to install keras 3 in a python venv, then try to run tests and see what breaks. It might be pretty bad.

It might also be worth asking Gemini. Maybe something like telling it to read the documentation for Keras 2, Keras 3, and and the TFQ documentation, then formulate a plan for how to go about updating the code base. Then again, there's a risk it will just blow through your token budget, due to the amount of text involved …

[rosspeili]

Thanks a lot @mhucka, I will do a comparison locally with local ollama models via continue plugin, so no token burning, but appreciate the heads up <3 I will share a short report here once I have a clear bird's eye view <3

[rosspeili]

@mhucka I did multiple sweeps with Gemini, Claude, and Composer APIs, as well as local DS coder and Qwen 3.5 coder, to ensure no angle is missed, and consolidated the report, although I have not run tests under Keras 3 ye, would be next step with local venv.

Summary

TFQ is Keras-2-only today through TF_USE_LEGACY_KERAS=1 (Bazel + CI + docs) and an explicit tf-keras dependency. The public API is built around 10 custom tf.keras.layers.Layer subclasses under tensorflow_quantum/python/layers/. They call TensorFlow Quantum custom ops and string circuit tensors in call(), so the realistic target for #1090 is Keras 3 with the TensorFlow backend, not JAX/PyTorch backends in a first pass.

Where legacy Keras is enforced

• configure.sh — TF_USE_LEGACY_KERAS=1 in Bazel env
• requirements.in / release/setup.py — tf-keras~=2.19.0
• scripts/ci_validate_tutorials.sh, scripts/test_tutorials.py
• docs/install.md + all 9 tutorial notebooks (“Configure Keras 2” cell)

Custom layers (migration surface)

Layer	Path
PQC, ControlledPQC, NoisyPQC, NoisyControlledPQC	python/layers/high_level/
Expectation, SampledExpectation, Sample, State, Unitary	python/layers/circuit_executors/
AddCircuit	python/layers/circuit_construction/elementary.py

Layers already use self.add_weight() for trainable parameters. Main migration themes from the Keras 3 migration guide:

1. Replace tf.keras imports with standalone keras
2. Expect to set jit_compile=False for models using TF string ops / TFQ custom ops (XLA default on GPU in K3)
3. Update loss/optimizer references in tests (eg. mean_squared_error call style)
4. Re-check multi-input functional models in executor tests (Keras 3 input nesting rules)

Example import change (pattern repeated across layer modules):

# Before (Keras 2 / legacy)
import tensorflow as tf
layer = tf.keras.layers.Layer

# After (Keras 3, TensorFlow backend)
import keras
layer = keras.layers.Layer

Example model compile tweak for TFQ custom-op layers:

model.compile(
    optimizer=keras.optimizers.Adam(),
    loss="mse",
    jit_compile=False,  # TF string tensors + tfq.* ops
)

Proposed smoke tests (next step)

Python 3.11 venv, TF 2.19.x, KERAS_BACKEND=tensorflow, without TF_USE_LEGACY_KERAS:

export KERAS_BACKEND=tensorflow
unset TF_USE_LEGACY_KERAS

bazel test //tensorflow_quantum/python/layers/high_level:pqc_test
bazel test //tensorflow_quantum/python/layers/circuit_executors:expectation_test
bazel test //tensorflow_quantum/python/layers/circuit_executors:state_test
bazel test //tensorflow_quantum/python/layers/circuit_executors:sample_test
bazel test //tensorflow_quantum/python/optimizers:spsa_minimizer_test

I will post pass/fail results in a follow-up comment after running these.

Suggested PR sequence

1. Layer code + layer unit tests (Keras 3 + TF backend)
2. Optimizer tests + remaining Python tests
3. Remove legacy Keras from configure.sh / CI once tests are green
4. Tutorials + install.md, and other complementary docs where Keras reference + dependency pins (tf-keras → Keras 3 alignment)

Out of scope for now

• JAX / PyTorch backends (TFQ ops are TensorFlow-specific)
• Model serialization / model.export() unless tests surface breakage

Question for you @mhucka

Do we want a transition period supporting both legacy Keras 2 and Keras 3, or a clean cut once CI is green? I can start testing once you confirm, or add to the scope. <3