Add training loop fusion

[seanmor5]

trafficstars

With the support of while loops in Nx and progress towards streaming data to the device, we should consider an API which allows for fusing the training loop into a single XLA computation. We will have to consider how this ties in with callbacks, hooks, etc.