Alexander Hsu

I ran into this same confusion for a little bit. How would you feel about adding something like a ScaledControlTerm, where the vector field always a multiple of the identity...

I see, that makes sense. There's currently no way to specify the diffusion in a matrix free way, right (other than WeaklyDiagonalControlTerm)?

That makes sense, thank you. This structure of vector field * control is really elegant, I've enjoyed working with diffrax and equinox a lot!

With the advent of [jax v0.6.0](https://github.com/jax-ml/jax/releases/tag/jax-v0.6.0), jax.tree_map has been removed and will error out now, so it'd be nice if this could be merged and have a new release?