Conversation
|
|
||
| ### Compilation is taking a long time. | ||
|
|
||
| - Set `dt0=<not None>`, e.g. `diffeqsolve(..., dt0=0.01)`. In contrast `dt0=None` will determine the initial step size automatically, but will increase compilation time. |
There was a problem hiding this comment.
Are the first two of these points worth keeping anywhere?
There was a problem hiding this comment.
They feel niche enough (i.e. not meeting the 'F' in 'FAQ') that I was feeling inclined to cut them.
If we ever find ourselves with a longer list of tips-and-tricks to tackle compilation time then I'd be happy to give them a home there, however.
| import timeit | ||
|
|
||
| @jax.jit | ||
| @eqx.debug.assert_max_traces(max_traces=1) |
There was a problem hiding this comment.
Does this have any overhead? I.e. should it be used to debug, then excluded from the final timing analysis (which is this code block)?
There was a problem hiding this comment.
Oh, good catch. Only overhead at compile time, I think! So we'll be fine for the use-case here.
| execution_time = min(timeit.repeat(lambda: jax.block_until_ready(run(x)), number=1, repeat=20)) | ||
| ``` | ||
|
|
||
| 2. Use the same ODE solver in both implementations to get an apples-to-apples comparison. It's not surprising that different solvers give different performance characteristics. (And if one implementation does not provide a solver that the other does, then no comparison can be made.) |
There was a problem hiding this comment.
"same ODE solver" -> "same solver"?
There was a problem hiding this comment.
So here (and for your comment below) I'm intentionally only focusing on ODE solvers.
The reason for this is that I already know that we could do more to improve performance of SDE solvers, and I've just never really found the time to figure that one out. (And on this note you've been very gracious around not merging the stateful-controls PR on this topic!)
| execution_time = min(timeit.repeat(lambda: jax.block_until_ready(run(x)), number=1, repeat=20)) | ||
| ``` | ||
|
|
||
| 2. Use the same ODE solver in both implementations to get an apples-to-apples comparison. It's not surprising that different solvers give different performance characteristics. (And if one implementation does not provide a solver that the other does, then no comparison can be made.) |
There was a problem hiding this comment.
I guess timing has been mostly ODE focused, but for SDEs is there anything of note (e.g. in how brownian noise is computed between implementations)?
|
No notes, this looks great! Very useful to be able to point people to this. |
|
Thank you both for the comments! |
3 participants
Fixes #688.
Tagging @lockwo @johannahaffner WDYT?
(Whilst I'm here I've removed the 'compilation is taking a long time' which is usually much less of an issue these days, ever since we landed the scan-over-stages refactor. I also don't like that it recommends
noinlinethese days, which is pretty deep magic!)