Follow-up from "Implement Dataclasses" concerning the shape

we recently merged quite a lot changes into main and #12 (closed).

One aspect remained unclosed, the following discussion from !7 (merged) should be addressed:

@hanikevi started a discussion: (+4 comments)

I don't quite get why synth_data_batch is constructed here, couldn't we just tree_map batch itself? Also; I'm not sure this works if we have, e.g. a 6-dof F/T signal, so the leaf is of size (batch_size, window_size, 6), right?