patrick-kidger / diffrax

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/
Apache License 2.0
1.45k stars 130 forks source link

Parallelize testing #487

Closed lockwo closed 2 months ago

lockwo commented 3 months ago

90-95 min is benchmark to beat (based on https://github.com/patrick-kidger/diffrax/pull/463, https://github.com/patrick-kidger/diffrax/pull/469, etc). Ideal is 4x since I believe default runner has 4 CPUs, but I don't think that will happen since its run one file at a time (maybe 2x speedup will be good?). No ooms locally, but not sure how much memory the runner has.

Addresses https://github.com/patrick-kidger/diffrax/issues/486