Viewing a single comment thread. View all comments

fhchl t1_izif55h wrote

Is this feature of torchode of solving multiple ODEs at once over some batch dimension comparable to jax.vmapping over that dimension in diffrax?

1

martenlienen OP t1_iziq1xp wrote

Yes, it is the same thing. Unfortunately, functorch is not advanced enough yet to just translate diffrax to PyTorch directly. Instead, we had to take care of batching everywhere explicitly to decide how long to loop etc.

1

fhchl t1_izjb0tf wrote

Aight! Thanks for the nice answers! I wish a good conference :)

1