[benikm91]

We have a quite complex, rarly used rearrange method that was designed very early when experimenting with the type system. While it is an impressive method, it is rarly used, quite complex in implementation, and requires einops dependency in the underlying Python environment. Due to its rare use we are also unsure if it works with jit, automatic differentiation, etc. There might be unexpected limits.

Due to these reasons and mostly due to the extra einops dependency this PR suggests removing the rearrange method for now. The code was backed up in an internal repo "dimwit-einops" that we might publish in the future. This follows our decision for dimwit-sharding, that we keep the core minimal while possibly adding more features with additional packages. So, if rearrange makes a comeback it will probably be as an additional package outside of dimwit (the core).

Overview changes:

• Remove rearrange: Remove operation. Type level stuff that was rearrange specific (specifically for the nice error message we had). Remove test. rearrange was not documented in AGENTS.md, so no removal in this documentation.
• VAE: We actually used rearrange in the VAE example, which now uses flatten and unflatten. I added an extra function to make the code still quite nice to read.
• Dockerfile: This PR also removes einops from the Dockerfile, which still installs, e.g., matplotlib. Let me know if we want to rollback on that and keep einops there as a dependency. I am quite unsure about it.

[marcelluethi]

Great change. It is telling that the changes to our code base are actually quite minimal after removing this metod.