Callbacks Quality of Life improvements

[lorenzoh]

Collecting some nice-to-have changes to callbacks. Some of these are breaking so they should be part of the next minor release.

• CustomCallback: Pass event and phase through to the wrapped function. Not having access to this information is unnecessarily limiting. BREAKING
• ToDevice: Change default behavior so that every array that is in state at StepBegin is moved. Add option to only move specific keys. This is a good default when implementing custom training Phases, as the current implementations assumes there are always xs and ys (this is from before the new training loop API was introduced). BREAKING