[jpf0]

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]

I 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]

Correct! Consider when you call "defn" for the first time to be the entry-point. Everything inside that is fused.