Currently, I'm currently using TF 1.x for most of my AD needs, but encountering the limits so I'm always looking for neater solutions. Conceptually, I enjoy th

Currently, I'm currently using TF 1.x for most of my AD needs, but encountering the limits so I'm always looking for neater solutions. Conceptually, I enjoy the functional approach where the standard compiler does the work. But I'm still using some graph based approach, where the order of operations are recorded on a "tape", then replayed it in reverse order. The main reason is to be able to manage the memory and computation placement. To distribute huge computations among machines, while limiting memory transfers between machines/devices to a minimum. Also some kind of automatic gradient accumulation to reduce the memory need would be nice. Currently using the graph based approach, partitioning a big tensor into smaller ones and doing a map-reduce to accumulate the gradient works surprising well (except for the initial very long graph creation time). Every time I do one of those manual optimizations I start dreaming : "that should be done by the compiler". You tell the compiler your cluster definition and devices, you write a few for loops as if the code was done on a single cpu. And the compiler do the loop splitting, reordering, parallelization across device in an optimal way given your available memory constraints.