patrick-kidger / diffrax

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/
Apache License 2.0
1.38k stars 124 forks source link

Improve errors #478

Closed lockwo closed 1 month ago

lockwo commented 1 month ago

Previously, the abstract term error, now:

https://github.com/patrick-kidger/diffrax/issues/446: ValueError: Error while tracing ODETerm(vector_field=<function equations>).vf: matmul input operand 1 must have ndim at least 1, but it has ndim 0

https://github.com/patrick-kidger/diffrax/issues/461: ValueError: The brownian increment BrownianIncrement(dt=f32[], W=f32[5]) does not have the minimal Levy Area <class 'diffrax._custom_types.AbstractSpaceTimeLevyArea'>.

https://github.com/patrick-kidger/diffrax/issues/474: ValueError: The brownian increment BrownianIncrement(dt=f32[], W=f32[]) does not have the minimal Levy Area <class 'diffrax._custom_types.AbstractSpaceTimeLevyArea'>.

Just generally propagating the actual error

patrick-kidger commented 1 month ago

(Let me know when you want a review btw, I'm ignoring this for now whilst it's a draft -- which I'm taking as a "you don't need to review this yet" signal)

lockwo commented 1 month ago

(Let me know when you want a review btw, I'm ignoring this for now whilst it's a draft -- which I'm taking as a "you don't need to review this yet" signal)

Yea, that's exactly what I mean by draft, you can ignore it. But with the weird backward edge case resolved, it is ready for you to take a look at

patrick-kidger commented 1 month ago

Great! This LGTM.

Looks like one of the tests has gone a bit wonky, I'm not sure why. Anyway, happy to merge this once things are green.

(Whilst I'm here, I'm also noticing that our tests have quite a few warnings in them. Some of these are expected -- fixed upstream in Equinox, which should have a new release any day now. The one about WeaklyDiagonalControlTerm could probably do with being suppressed in the tests, though.)

patrick-kidger commented 1 month ago

And merged! :)