KeyError: dtype([('float0', 'V')]) grad + odeint in JAX 0.3.3 #9951
-
|
Hi all, I've just upgraded to 0.3.3 and am getting a new error I've never seen before in some code that executed correctly in previous versions. I don't yet have a minimal example, but wanted to check in here if anything rings a bell, as the error message is pretty uninformative (to me). The exception being raised is extremely long, but it is coming up when trying to take the |
Beta Was this translation helpful? Give feedback.
Replies: 7 comments 32 replies
-
|
Thanks for raising this. I have a guess... but is there a chance you could share a runnable repro so I can test my guess? |
Beta Was this translation helpful? Give feedback.
-
No worries! Hopefully the above example is useful in figuring out what's going on. |
Beta Was this translation helpful? Give feedback.
-
|
That is suuuuuper useful, thank you! |
Beta Was this translation helpful? Give feedback.
-
|
This is great. One issue is that The other issue is that we shouldn't have been letting float0 arrays propagate so far. There was a typo in #9923! Luckily with your repro it was super quick to find. I'll push a PR with a fix momentarily. Maybe we don't need to roll back #9923 after all... However, that still leaves the slowdown that @oracle3001 observed, so we'll need to check if this typo is related to that slowdown. |
Beta Was this translation helpful? Give feedback.
-
|
#9959 has the typo fix. I may add your example as a regression test for it, unless I can put more thought into how to reproduce it more simply. |
Beta Was this translation helpful? Give feedback.
-
|
Awesome! I just checked out #9959 locally and confirmed my original test case is working again. |
Beta Was this translation helpful? Give feedback.
-
|
Has there been some significant change to the odeint method (or changes that directly affect its running) between v0.3.1 and v0.3.3, as some code I have now runs 50% slower? |
Beta Was this translation helpful? Give feedback.
-
|
Thanks for checking that. Let's roll back #9923 and do a release... |
Beta Was this translation helpful? Give feedback.
-
|
@oracle3001 could you patch #9959 and see if the slowdown is still present? |
Beta Was this translation helpful? Give feedback.
-
|
After #9923 , optimization with |
Beta Was this translation helpful? Give feedback.
-
|
Worth noting that the slowdown may actually be desirable for correctness. Previously |
Beta Was this translation helpful? Give feedback.
-
|
@YouJiacheng in #9923 due to a simple bug ("typo") we were hoisting out integer valued arrays, i.e. ones that we didn't need to hoist out. So in effect it overcorrected from the previous bug where some were being missed. The fix in #9959 tightens things up. |
Beta Was this translation helpful? Give feedback.
-
|
(Starting a new unified top-level thread.) FYI @oracle3001 @DanPuzzuoli we pushed a new pypi version of jax==0.3.4 with the #9959 fix merged in. |
Beta Was this translation helpful? Give feedback.
-
|
Thanks for the quick solution! |
Beta Was this translation helpful? Give feedback.
-
|
With v0.3.4, I am still suffering the extreme slow down. It might be as Patrick says due to a bug in the previous code missing inputs, but it certainly an extremely significant slow down. |
Beta Was this translation helpful? Give feedback.
-
|
How does the speed change when switching out For context, |
Beta Was this translation helpful? Give feedback.
-
|
Diffrax is 2x slower than even odeint in v0.3.4. However with the sort of problems I am working with even simpler examples Diffrax is normally noticeably slower (with the same solver, rtol, etc) than odeint. |
Beta Was this translation helpful? Give feedback.
-
|
Can you open a MWE as an issue on the Diffrax GitHub page? Diffrax is usually 10-50% faster in my experience, so I'd like to investigate this. |
Beta Was this translation helpful? Give feedback.
-
|
I'm getting a very similar error message, and this discussion was the only hit on google when searching. But I've got jax 0.3.13 installed, which has the fix discussed above. For anyone else encountering this, the problem is caused by numpy 1.23.0 and can be worked around by downgrading numpy. See #11221 |
Beta Was this translation helpful? Give feedback.
-
|
Hello everyone, The output is the following: I am using |
Beta Was this translation helpful? Give feedback.
-
|
This looks like the bug reported in #11221 , which was fixed in jax version 0.3.14. I see that you report using jax v0.4.1, but please double-check that: the line numbers in the traceback you report indicate that you may be using JAX v0.3.13 or older. Compare these:
The fix to the issue you're seeing is either to update JAX to |
Beta Was this translation helpful? Give feedback.
-
|
Hi all, I am getting a similar error as cited above while differentiation through an odeint and linalg.eigvals. I am using JAX 0.4.14 and numpy 1.23.1. Thanks for your help! |
Beta Was this translation helpful? Give feedback.
-
|
Maybe try Diffrax instead? |
Beta Was this translation helpful? Give feedback.
-
|
Can you include a minimal reproducible example of the error you're seeing? It's hard to guess what's going on from a traceback alone. |
Beta Was this translation helpful? Give feedback.
-
|
Hi everyone, thanks for your replies. I solved the issue, somehow. Thanks again. |
Beta Was this translation helpful? Give feedback.
-
|
Hi @pragati903 how did you solve the problem? Having the same issue with jax 0.4.13 and numpy 1.23 and downgrading numpy to <= 1.22 didn't solve the issue. Weird is my code used to run fine until just a few months ago. |
Beta Was this translation helpful? Give feedback.
-
|
Hi, sorry you're having issues! If you can share a minimal reproducible example, that would help. I'd consider doing so in a new issue – an answer on a year old discussion thread is not the easiest place for us to track new questions or issues. |
Beta Was this translation helpful? Give feedback.
-
|
Thanks for the reply. Of course, though the code resides on a virtual machine and reads private data. I will try to create a reproducible example with fake but similar enough data that I can share here. |
Beta Was this translation helpful? Give feedback.
-
|
Hello, I am getting this error again on Python 3.10, Jax 4.28: The issue in my case is when one of the *args passed to func in odeint is an integer. For some reason, this causes the above error when backproping through the odeint. In my case, I am able to get around this error by passing the arg as a float, and then casting it to int inside of func using Here is the docstring for odeint for reference: |
Beta Was this translation helpful? Give feedback.
-
|
Thanks for the report! If you're able to put together a minimal reproducible example of the issue you're seeing, we'll try to fix it. |
Beta Was this translation helpful? Give feedback.
Thanks for raising this. I have a guess... but is there a chance you could share a runnable repro so I can test my guess?