Would Diffrax provide any value as a replacement to torchsde?

[polarathene]

From: google-research/torchsde#128 (comment)

One thing you could try is to switch to Diffrax. By taking advantage of JAX's JIT compiler, then this can sometimes be substantially faster than torchsde. (Probably due to reduced Python overhead and fewer memory allocations.)

The linked issue is regarding training, but I came across it while looking into the torchsde dependency. No idea if adopting diffrax would provide any additional value to how torchsde is used within ComfyUI, but thought I'd share for awareness 😅