Implement Dataclasses

Hey @hanikevi,

We have some open questions we want do discuss with you:

• what about the collate_fn: in the current implementation we use a fn that casts dataclasses into jax arrays. So we are not sure whether we need some modifications to train with dataclasses in the ae,
• where to put the iter fn: think about moving the iter func into sysvars (actual iter func is in synth dataclass).

Some Todos were added in the code.

Closes #12 (closed)