9 ms·
No autodiff for most of these though.
by hyperbovine 6y ago
No autodiff for most of these though.
- alevskaya 6y agoJAX autograd will work on most any jitted fn - the control-flow limitations are no autograd for code with for/while loops since there's a statically unknowable trip count through the loop body. Much looping code can be handled differentiably using a "scan" though.