Report: Behavior of external variables with TensorFlow2.x backend

[KentaroAoyama]

Dear Developers,

I have so much respect for your hard work.
I have found a problem with the behavior of using external trainable variables in the Tensorflow 2.x backend and would like to report it. I hope it will be of some help to you.
The version of DeepXDE is 1.12.0, and the version of TensorFlow is 2.17.0 in my environment.

The bug was caused by using the example described in the DeepXDE document that is modified to use external trainable variables as a TensorFlow 2.x backend.
For example, an error occurred when adding the codes below to this example's head.

from os import environ
environ[“DDE_BACKEND”] = “tensorflow”

The above code throws an error: AttributeError: 'ResourceVariable' object has no attribute 'overwrite_with_gradient.'
This error is because dde.Variable is a tf.Variable (ResourceVariable) without an "overwrite_with_gradient" attribute value, while the neural network model parameters are implemented in keras.src.backend.sensorflow.core.Variable, which has an "overwrite_with_gradient" attribute value.
After rewriting dde.Variable to keras.src.backend.tensorflow.core.Variable, the above error no longer occurs, and the training of the external trainable variable is progressing.

I would like to suggest the following solutions.
(1) Implement dde.Variable in keras.src.backend.tensorflow.core.Variable (I don't know the extent of the impact, this could be dangerous)
(2) Replace dde.Variable with keras.src.backend.tensorflow.core.Variable at compile stage.
(3) Do not change the implementation, but specify the above specifications in the docstring and example codes.

Thank you for taking the time to read this.

[lululxvi]

Thank you for reporting this. This code should work well with older versions of TensorFlow. So basically what you mean is that the code failed with TensorFlow 2.17.0 and "tensorflow 2.x" as the backend. I guess this is because TensorFlow 2.16+ uses Keras 3, which introduces many issues. It would be great if you could confirm this. We can open a new PR to fix the code.