Class based recorders

[SamTov]

While writing some recorders copilot suggested the following API for adding a class-based recorder:

for i in range(10):
    train_recorder.add_class_recorder(i)

Seems pretty good in the end... You can do the partitioning of the train ds inside of the current recorder and either add a group or spawn a new recorder using data only from that class.

[KonstiNik]

I already did a draft at some point for that using vmap, which essentially is the same. It is nearly done.
1cf02d9
Is that similar to your idea?