Dynamic Automatic Differentiation of GPU Broadcast Kernels · HackerPapers discussions

Dynamic Automatic Differentiation of GPU Broadcast Kernels — mentioned by jrevels

Tue, 18 Jun 2024 18:20:30 GMT

Cool stuff! Thanks for the write-up. Getting broadcast ergonomics right in a fully statically typed context seems both tricky and pretty satisfying. > We thought it’d be a fun/simple feature and a good break from the heavy duty work on automatic differentiation we had just published Perhaps it wasn't as much of a break as you might think - funnily enough, static broadcast semantics can serve as an interesting structural foothold for certain AD optimizations [1] (disclaimer: I'm a coauthor on that paper :) but it's been quite a while since I've done work in the AD space) [1]…

Paper and discussion history. A mention is not an endorsement.

Dynamic Automatic Differentiation of GPU Broadcast Kernels — mentioned by mattjjatgoogle

Mon, 16 Mar 2020 19:32:36 GMT

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://openreview.net/forum?id=BkYYXJ9i-). I'm hoping as users push JAX on new applications, especially in simulation and scientific computing, we'll learn a lot and be able to improve! There's also "cross-country optimization" (https://www-sop.inria.fr/tropics/slides/EdfCea05.pdf) for mixing some forward-mode into reverse-mode to improve memory efficiency. Analogously to jax.checkpoint, we've only experimented with exposing…

Paper and discussion history. A mention is not an endorsement.

Dynamic Automatic Differentiation of GPU Broadcast Kernels — mentioned by jrevels

Thu, 08 Nov 2018 05:21:39 GMT

Author here; the arxiv version can be found at https://arxiv.org/abs/1810.08297. Not much different from OP's linked version, but it includes citations to other interesting Julia AD/TPU-related papers that utilize this technique. Happy to answer any questions, at least until I turn in for the night :)

Paper and discussion history. A mention is not an endorsement.

Dynamic Automatic Differentiation of GPU Broadcast Kernels — HN submission

Tue, 23 Oct 2018 18:57:14 GMT

Dynamic Automatic Differentiation of GPU Broadcast Kernels

5 points; 1 discussion reply on HN. Paper and source history. A submission is not an endorsement.