jit and vamp best practices

patrick-kidger / quax

Multiple dispatch over abstract array types in JAX.

Apache License 2.0

106 stars 3 forks source link

Open nstarman opened 9 months ago

nstarman commented 9 months ago

Just a suggestion to add a small section to the docs about best practices for quaxifying jitted and vmapped functions. Thanks!

patrick-kidger commented 9 months ago

At least with JIT, ths usual best-practice of putting JIT at the top level applies.

For vmap, I don't think it matters: put these in whichever order describes the operations you're trying to perform.

Does that sound reasonable to you?

(I agree it'd be worth emphasising this somewhere.)