`if` Statements Behave Differently in `jax==0.4.1`

LouisDesdoigts / dLux

Differentiable optical models as parameterised neural networks in Jax using Zodiax

https://louisdesdoigts.github.io/dLux/

BSD 3-Clause "New" or "Revised" License

43 stars 6 forks source link

`if` Statements Behave Differently in `jax==0.4.1` #182

Closed Jordan-Dennis closed 1 year ago

Jordan-Dennis commented 1 year ago

Hi all, if statements no longer seem to be compilable (at all) in the latest jax. I think any run-time if statements (there aren't many) need to be moved over to cond. Regards Jordan

LouisDesdoigts commented 1 year ago

I can't see anything in the changelog, what exactly is the different behavior? Using cond with python booleans breaks some things so id be surprised if this was the case

Jordan-Dennis commented 1 year ago

Wasn't sure the issue popped up and disappeared. I don't know what caused it. I think it might have been something persisting in the ipykernel, that wasn't meant to be there.