Solving simple dynamics: ControlTerm piecewise product

Hi,
Thanks first for developing this nice package.

For the context, I intend to use diffrax to implement a custom Langevin-like dynamic, but my issue can be reduced to the following. Let's say I want to implement a simple $n$-dimensional Brownian motion:
$$dX = dB$$

I can try doing

diffusion = lambda t, y, args: jnp.ones(n)
brownian_motion = VirtualBrownianTree(t0, t1, tol=dt/2, shape=(n,), key=seed)
terms = ControlTerm(diffusion, brownian_motion)

but it wouldn't work because of this line defining the vf-contr product for ControlTerms. What tensordot(vf, contr, axes=ndim(contr)) does is fully contracting tensors (on all the dimensions of contr), so in my case it would return a scalar $n dB$, whereas I would generally require piecewise (Hadamard) product vf * contr.

For now, the only way I found to implement piecewise ControlTerm product is to increase the dimensionality of the vector field, e.g. in that case, write diffusion = lambda t, y, args: jnp.eye(n), which is way more expensive ( $O(n^2)$ ) and will not scale to my applications. And I am not sure some jax.experimental.sparse matrices would help.

I understand matrix product, nay higher rank tensor products, may be required in some applications. This recent question, or this diffrax example of Neural SDE, have both matrix-valued diffusion vector field and vector-valued Brownian control. However, if I am not wrong, it seems that for that same reason of full tensor contraction, matrix product between matrix-valued vf and matrix-valued control is not currently easily implemented.

So my question would be: Did I miss a way to implement ControlTerm piecewise product?
I think it should be possible to implement it:

• without much additional computational cost.
• by conserving PyTree structures (without having to tree flatten everything manually) as diffrax already handles it well.
• and optionally, by allowing other tensor contractions such as matrix-matrix matrix product.

I could not think of any einsum to replace tensordot(a,b,ndim(b)) that would fit well in all cases, but maybe having a way to specify which product _prod function to use in ControlTerm could be an idea? Or maybe I just missed a simple way to do everything above.

Thanks in advance!

Start with diffrax/_term.py around line 250 and inspect ControlTerm's vector-field/control product, using the Brownian-motion example in the issue as the first reproduction. Compare the current full contraction with the requested piecewise product and consider how existing PyTree handling and other tensor contractions are represented. Done means the desired product behavior is supported without the unnecessary identity matrix workaround.