Closed AbhiDu96 closed 11 months ago
Thanks. With jax_enable_x64, the argmax output was int64. I changed the code to use the expected int32.
The implementation is not tested with jax_enable_x64. Check that your results are sensible.
Hello @fidlej , Thanks for the solution. With the corresponding changes it works without any warnings or errors. Although I still have to check if the results with tax_enable_x64 are sensible. I think we can close the issue with this. Thanks again.
Hello, Thank you for this amazing library. It's been a while I am using this package for our project. I have started facing the following issue: ''' /home/dubey/anaconda3/envs/jaxenv_qiskit/lib/python3.11/site-packages/jax/_src/ops/scatter.py:93: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=int64 to dtype=int32. In future JAX releases this will result in an error. warnings.warn("scatter inputs have incompatible types: cannot safely cast " ''' I want increased precision for my experiments and so I set the following global variable as: ''' jax.config.update("jax_enable_x64", True) ''' Unfortunately, I have tracked down the bug to gumbel_mu_zero_policy function and within these functions the
I am not sure how to resolve this issue but kindly let me know if this does not make sense.
Update: I tried keeping "int32" everywhere but as expected that does not make sense since when you specify the global config to x64 it does change precisions of everything.