Closed amunozj closed 5 months ago
In general, this sometimes happens when a rouge operation creates a float32 when float64 is enabled. This could happen for instance if the below is not at the start of the python program.
from jax import config
config.update("jax_enable_x64", True)
Can you confirm the location of the above?
If JAXNS code loads before this then internal dtypes can be set to 32-bit after which enabling 64-bit would allow constant conversion to create mis-matches. In general this can probably be resolved by carefully controlling the input/output dtype invariants but that will take some work and testing.
Thank you so much @Joshuaalbert ! making sure the definition of the x64 config is set, as early as possible, does solve the issue!
Describe the bug I'm using NumPyro's JaxNS wrapper. In the forum I asked about running JaxNS in 64 bit. See #153 In that discussion the recommendation was to set the 64-bit flags. I'm currently setting:
Expected behavior I expected the sampler to run in 64-bit
Observed behavior The sampler crashes because of a type mismatch:
The problem seems to be with these two variables:
Removing the 64-bit flags removes the crash. Any suggestions as to how to fix?
JAXNS version jaxns==2.4.10