Viewing profile — mattjjatgoogle

An author's tweet thread: https://x.com/jacobaustin132/status/1886844716446007300

Actually that never changed. The README has always had an example of differentiating through native Python control flow: https://github.com/google/jax/commit/948a8db0adf233f333f3e5…

Actually, that's never been a constraint for JAX autodiff. JAX grew out of the original Autograd ( https://github.com/hips/autograd ), so differentiating through Python control flo…

You're right! Maybe we should revise that... I made https://github.com/google/jax/pull/17851 , comments welcome!

Have you seen JAX MD? https://github.com/jax-md/jax-md

You're right that downstream libraries have often tended to introduce magic (some more than others), and moreover one library's magic is typically incompatible with other libraries…

Thanks for taking the time to explain these. > It's been a bit, but I think the most frustrating errors were around mapping pytrees (like this issue https://github.com/google/jax/i…

If you have any particular examples in mind, and time to share them on https://github.com/google/jax/issues , we'd love to try to improve them. Improving error messages is a priori…

Thanks for the kind words! We've been doing a lot more work in this direction too, for both compiler-based automatic parallelization [0] and a work-in-progress pmap successor for '…

I'm sure there's a lot of good material around, but here are some links that are conceptually very close to the linked Autodidax. (Disclaimer: I wrote Autodidax and some of these o…

The difference is that in TF1 you had to use tf.cond, tf.while_loop etc for differentiable control flow. In JAX you can differentiate Python control flow directly, e.g.: In [1]: fr…

No worries! I didn't mean it as a correction so much as just a discussion; I'm sure it's true that other autodiff systems have very sophisticated automatic remat (like https://open…

Rematerialization in autodiff is super interesting! XLA does rematerialization optimizations, so you get those automatically under jax.jit. There's also the jax.checkpoint decorato…