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