vector-valued SDE cumbersome workflow

[lostdevfound]

It seems the vector-valued SDE workflow is very cumbersome and I couldn't get the module working past scalar SDEs. Here is the minimal example of a vector-valued SDE. Whatever I try, there is always some issue with batch_size or matrix dimensions.

k = 0.5
w = 1
Z = torch.tensor([[1, 0],[0, -1]])
I = torch.eye(2)
H = w*Z

class SDE(nn.Module):
    def __init__(self):
        super().__init__()
        self.noise_type = "diagonal"
        self.sde_type = "ito"
    
    def f(self, t, y):
        return (-1j * H -2*k*I) @ y
    
    def g(self, t, y):
        return I

batch_size, state_size, t_size = 1, 2, 100
sde = SDE()
ts = torch.linspace(0, 1, t_size)
y0 = torch.full(size=(batch_size, state_size), fill_value=1 + 0j)

with torch.no_grad():
    ys = torchsde.sdeint(sde, y0, ts, method='euler')

[patrick-kidger]

If you are using noise_type = "diagonal" then the diffusion should return the diagonal as a vector, rather than as a matrix. That's probably the issue.

FWIW I would recommend considering Diffrax instead, which is much faster, still under active development, and more feature-complete. (At this point I think it's fair to say that torchsde is just in maintenace mode.)