Checkpointing arbitrary tags

[connorjward]

As discussed in #3310 (review), it should be feasible to design a modified interface to save_function such that arbitrary tags can be passed. I am suggesting this generalisation because changes are currently being proposed that would add 2 additional timestepping arguments to save_function (when one, idx, already exists). This leads to code duplication and a more complicated API. @ksagiyam and I feel like it should be possible to instead have an API like: save_function(func, **kwargs) where **kwargs represents arbitrary information to be stored.