custom_gradient not working with JAX backend

[larschristensen]

When trying to use the custom_gradient decorator with a JAX backend as shown in the example located here
custom_gradient_example, I get an error message saying TypeError: 'NoneType' object is not callable.

It works without a problem on the TensorFlow backend, but not JAX. The behaviour is seen both on keras v3.9.0 and on keras-nightly.

[dhantule]

Hi @larschristensen, Thanks for reporting this.

I am unable to reproduce your error. I have tested the custom_gradient example with Keras 3.9.1 and I'm able to run it with JAX backend without any issues in this gist.

[larschristensen]

Hi @larschristensen, Thanks for reporting this.

I am unable to reproduce your error. I have tested the custom_gradient example with Keras 3.9.1 and I'm able to run it with JAX backend without any issues in this gist.

Hi @dhantule. I also don't see the problem until the function is used together with an optimizer. Try to use the function in a small neural net and optimize it, so that the custom gradient is utilized by the optimizer.

[dhantule]

Hi @larschristensen, could you provide some code to reproduce ? Thanks!

[larschristensen]

I thought it would be relatively easy to recreate this failure in a simple example, but after having tried to do so, I have not been able to.

The problem I'm seeing in more complicated code, which I'm unable to share, is that the line https://github.com/keras-team/keras/blob/master/keras/src/backend/common/variables.py#L265 is being hit with self._value being None when using the JAX backend. As self._initializer is also None, it fails with the above mentioned error message. Looking at the stack trace, the call that fails originates from some JAX XLA code, so the problem may in fact lie with JAX / XLA and not Keras. When using the TensorFlow backend, this doesn't happen as the variables then have a numeric value and everything works as expected.

I have worked around the problem by replacing custom_gradient with stop_gradient, so I think we can close the issue for now. However, there seems to be a problem with using custom_gradient on the JAX backend that can be triggered by some but not all code. Hence, the underlying issue on the JAX backend remains unresolved.

[larschristensen]

Hi @dhantule

I had another look at this and finally managed to create a modified version of timeseries_weather_forecasting.py which reproduce the issue. You can find the code here: https://colab.research.google.com/drive/1514lSXjRH7pIjyEdJSohXZSURFX-kic6?usp=sharing. Running the code as-is using a JAX backend should show the problem.

Could you take a look at it?

[JyotinderSingh]

Please provide the reproduced issue in the form of a colab notebook, this ensures that the issue is easily reproducible on standard hardware.

Also, maintainers won't be downloading archives from the internet due to the inherent security risks associated with that.