bug: Check dataset shapes when adding.

[daniel-dodd]

The __add__ method that concatenates two jaxutils.Dataset's does not provide any checks for both dataset batch shapes. Though this will error via jnp.concatenate, it would be nice to write a function that checks shapes and gives a clear error message to users.