IID Brownian motion?

Hey there,
First off, thanks a lot for this library!
I wondered whether it is possible to solve a 2-dimensional SDE with IID noise? As a simple example, consider the Ornstein-Uhlenback process (I need it to drive other ODE, so I want to use sol_ou.evaluate(t) in the later vector field).

t0, t1 = 0.0, 20.0
DT = 0.001
dim = 2

def drift_ou(t, y, args):
    return (args["mu"] - y) / args["tau"]

def diffustion_ou(t, y, args):
    return jnp.ones((dim))*args["sigma"]

brownian_motion = dfx.VirtualBrownianTree(
    t0,
    t1,
    tol=1e-3,
    shape=jax.ShapeDtypeStruct((dim,), np.float64),
    key=jr.PRNGKey(1),
)

terms = dfx.MultiTerm(dfx.ODETerm(drift_ou), dfx.ControlTerm(diffustion_ou, brownian_motion))
args = {"mu": 1.0, "tau": 0.1, "sigma": 0.3}
sol_ou = dfx.diffeqsolve(
    terms,
    dfx.Heun(),
    t0,
    t1,
    dt0=DT,
    y0=jnp.ones((dim))*args["mu"],
    saveat=dfx.SaveAt(dense=True),
    max_steps=int(t1*(1./DT)),
    args=args
)

In this code, I am trying to create a 2-dimensional Ornstein-Uhlenback, and VirtualBrownianTree gives me the dW in 2 dimensions; however, the Brownian motion is the same for both dimensions. Is there some way to make the "samples" independent?

Cheers!

Start by reading the VirtualBrownianTree and ControlTerm usage in the issue example, focusing on how a vector-valued Brownian path supplies increments to the diffusion term. Use the two-dimensional Ornstein-Uhlenbeck example to determine whether independent coordinate samples can be supported, and verify that the resulting dimensions are no longer identical.