Memory safe computations with XLA compiler - Artemev, Artem; An, Yuze; Roeder, Tilman; van der Wilk, Mark (Advances in Neural Information Processing Systems, 2022)