AustinT/rough-adkf-jax

Better support for batching

Opened this issue · 0 comments

Is there a way to easily do batching across multiple tasks? Not sure if vmap could be used.