I greatly appreciate your work, by the way, I just have not been able to make use of it effectively. Some of the problems revolve around the question of how to best proceed in a situation, where you have a high dimensional state space, but it is naturally partitioned in some way. There are several solutions for this (SubArray etc.), but the burden is entirely put on the user. I had the impression that there was a tension between what I would have preferred to write and what I could
easily write as soon as I attempted to generalize from the examples that I could find. With JAX the corresponding libraries operate on Pytrees and as a user you can specify equations pretty naturally and easily without much fuss. If you want to use XLA, MPI, CUDA in Julia it typically is the same, in theory it should be possible to make things work in practice I have struggled tremendously to do anything productive.
When I looked at the adjoint event handling code last, I couldn’t figure out in the implementation, whether the general case was handled correctly, especially since parts of it still seemed in flux. Writing similar code in JAX leaves close to no room for interpretation that the code is correct. I am sure most of it is down to familiarity. But since ultimately I want to do ML relared things, right now JAX and related libraries ties up things that are there much more neatly even though overall SciML implements a more comprehensive set of techniques. I am still closely following the work around it especially in the area I am interested in and have some prototypes written in it, but it just hasn’t clicked yet.