May not answer all your questions, but this may be a good starting point: https://hexdocs.pm/nx/Nx.Defn.html