HajimeKawahara / exojax

🐈 Automatic differentiable spectrum modeling of exoplanets/brown dwarfs using JAX, compatible with NumPyro and JAXopt
http://secondearths.sakura.ne.jp/exojax/
MIT License
45 stars 14 forks source link

Change the order of calculations for 32-bit mode #482

Closed sh-tada closed 5 months ago

sh-tada commented 5 months ago

I found that we can prevent NaNs in gradients that occur only in 32-bit mode by modifying the order of calculations. This may be because it is better to perform computation with 64-bit variables before applying 32-bit variables to minimize rounding errors.

Fixes #473