4 ms·
If I understand correctly, Nx.defn effectively calls a backend compiler for that function, and that will presumably result in only one round-trip of data to the
by jpf0 6y ago
If I understand correctly, Nx.defn effectively calls a backend compiler for that function, and that will presumably result in only one round-trip of data to the accelerator hardware (e.g. GPU). If you have two defn functions, do you 1) lose compilation between them and 2) round trip twice, or can you nest defn functions and get compilation of the composed function?
- dnautics 6y agoI believe that your understanding is correct: defn marks potential boundaries between "jit compiled" or even AOT compiled, and if you call a defn in a defn it will be fused.
- josevalim 6y agoCorrect! Consider when you call "defn" for the first time to be the entry-point. Everything inside that is fused.