Awesome! Looking at the definition for Axon.compile you’ll understand better that what your doing is basically inserting a few lines of code between the inner call for Axon.build (which yields non-compiled functions) and the actual Nx.Defn.compile call.
Before you were trying to feed an already-compiled function to the compiler yet again. But because it was already compiled, it impeded the compilation for your outer function.
