Extracting number of function and Jacobian evaluations from diffeqsolve

[amir-saadat]

Hey Patrick, Is it possible to get the number of function and Jacobian evaluations from the output of diffrax.diffeqsolve? Thanks

[patrick-kidger]

These aren't currently available but agreed that this would make sense as a statistic to collect.

This is the kind of thing that'd be really easy to do once (or if) stateful/effectful operations become a thing in JAX. And a huge pain to manually thread the counting operations through without it! [CC @LenaMartens for interest since we were discussing stateful operations recently.]

[LenaMartens]

Is this for logging purposes or would you use the value itself in a subsequent computation? Would a way to print these counts through eg. a JAX print be enough? (CC: @sharadmv)

[amir-saadat]

Thanks @patrick-kidger for the info. @LenaMartens, we'd like to log this ideally.