Maybe someone should extract the program flow from the AST to create the compute graph, or create a restricted language that facilitate that. It seems in TF or PyTorch the syntactic sugar is infective at some point and you have to use the explicit API.
I think that a data flow language would be interesting . Having dataflow as a first class values would be cool. Borrowing stuff from graph computation libraries or graph rewrite systems would be useful. The idea would be that you would take the individual pieces of a dataflow that could be packaged and manipulated . To give an example a set of layers of a CNN or even a whole network like inception could be a first class object. Also, you would want the dataflows to be able to be differential. I think some category theory properties would be helpful to use in this context .
Isn’t this what JITs do?
There's a (very exciting) PyTorch JIT incoming sometime soon!