Open anas-rz opened 11 months ago
So you've verified the StochasticDepth
layer works fine in JAX otherwise?
Typically the "tracer leak" type of issue happens when you take a tensor from an intermediate computation and store it as an attribute of a permanent object, like a layer or model.
I think the line random_tensor = keep_prob + keras.random.uniform(shape, 0, 1)
in the trace is interesting. What happens if you replace it with a constant, like 0.6?
@anas-rz , Could you please check with the latest Keras 3 package and let us know if you're still facing the issue.
If you could provide some sample reproducible code, it makes it easier for us to debug.
While converting a {model} from TensorFlow {to Keras Core}, I am facing error with the JAX backend at the end of the epoch. The error is thrown by the
StochasticDepth
layer {used here in the network}. The code works with PyTorch and TensorFlow backends and shows the leakage with JAX. I do understand the concept of pure functions in JAX. Still not able to figure out whether it's the problem of my framework or Keras Core. The StochasticDepth layer works perfectly CCT example {here}. Here's the stack trace if that can help: