Tensorflow 2.14 breaks dictionary output from a functional model

[mattdangerw]

When tf 2.14 releases, we will face a bug with dictionary output from a functional model.

https://colab.research.google.com/gist/mattdangerw/019d1ebcec746d6f7424248911cf16ae/tf-2-14-bug.ipynb

I am not sure this will show up on a colab right away, as it actually depends on the specific version of typing-extensions (this is a pretty obtuse bug). But we will see this in a clean pip environment which picks up the latest typing-extensions.

After the 2.14 release:

• Clean install keras-core and tensorflow in a pip environment.
• Attempt to write a functional model with dictionary output.
• Any compiled functions (e.g. predict, fit), will fail.

[mattdangerw]

Here's a bug about the root cause on the tensorflow side -> tensorflow/tensorflow#60687

Ideally this would just be fixed in tensorflow proper. On our side, we could do two things

• Avoid tracking out functional inputs and outputs. -> Hacky fix for dictionary output with tf 2.14 keras-core#933
• Upper bound typing_extension to 4.5, though this may conflict with other libraries.

[mattdangerw]

Roughly it looks like this is an issue with AutoTrackable tensorflow internals and the wrapt library tf uses for some of the tracking stuff.

[fchollet]

Thanks for the report. Is this being actively fixed by the TensorFlow team? Seems like a significant issue.

[mattdangerw]

Thanks for the report. Is this being actively fixed by the TensorFlow team? Seems like a significant issue.

No idea, commented on the bug and we can try to follow up more. It actually doesn't pop up with tf.keras, because tf.keras does not attempt to auto track functional tracing in the same way we are currently on keras-core. It only pops up when returning a dict from a tf.function that is a property of tf.Module, so it might not affect that much code.

[sachinprasadhs]

I think we can close this issue since it has no problem with the Keras 3 with the below code.

import keras

batch = keras.ops.ones((2, 1))
inputs = keras.Input(shape=(1,))
outputs = {
    "sum": keras.layers.Add()((inputs, inputs)),
    "diff": keras.layers.Subtract()((inputs, inputs)),
}
model = keras.Model(inputs, outputs)
model.predict(batch)

[github-actions]

This issue is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you.

[github-actions]

This issue was closed because it has been inactive for 28 days. Please reopen if you'd like to work on this further.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No