Code errors after upgrading from 0.4.29 to 0.4.30, help figuring out why?

[Chulabhaya]

Hi there! I recently upgraded my JAX version from 0.4.29 to 0.4.30, and all of a sudden I am receiving this error during the compilation I'm guessing? Here is the error: https://pastebin.com/0QQnyCzs

Can anyone help me understand what this error is saying and what could cause it to occur when upgrading from JAX 0.4.29 to 0.4.30? I tested several versions below 0.4.29 as well, such as 0.4.28, 0.4.27, and 0.4.26 and this error definitively occurs when I do the upgrade from 0.4.29 to 0.4.30. I went through the change log listed here for 0.4.30 and I didn't seem to spot anything that was a breaking change: https://github.com/google/jax/blob/main/CHANGELOG.md#jax-0430-june-18-2024

I can provide the actual code if requested, but I was first hoping to understand what this error itself means.

[yashk2810]

Can you upgrade to 0.4.31 and see if the error goes away?

[Chulabhaya]

Hi @yashk2810 , thank you for following up so quickly! I tried upgrading to 0.4.31 and I receive the same error. I am currently attempting to go through my code line by line with jax.config.update("jax_disable_jit", True), but would that actually help if this is an error that occurs during jit compilation somewhere? I don't quite understand the error message itself at the moment actually.

[yashk2810]

Do you have a minimal repro?

[Chulabhaya]

I don't currently because I have not been able to figure out what triggers the error, or where it is occurring. Do you have any tips as to how to narrow that down to be able to produce a minimal reproduction?

[Chulabhaya]

@yashk2810 So I've been slowly able to narrow down that the error is being caused as a result of this function, particularly line 43 of the loss. If I set that to a fixed value like 1.0, then the code will run; if I keep it as it is now, it will error out (but only on 0.4.30 and above).

https://pastebin.com/RHwud3hX

That loss function is being used for calculating the gradient:

(loss, qvals), grads = jax.value_and_grad(_loss_fn, has_aux=True)( train_state.params )

Do you see anything in there where there was something that changed between 0.4.29 and 0.4.30?

[Chulabhaya]

Is there any chance some sort of Jax caching or force quitting during a training can result in this error? I was temporarily able to fix this issue by updating my version of Flax to 0.9.0 and was able to run the training as normal (with Jax version 0.4.31). However after I force quit from that training and attempted to re-run, it immediately crashed with that error.