Clarification on 3D+ / Convolutional tensor support (`flatten=True`) in Dion

[JenWei0312]

Hi team,

Thank you for open-sourcing Dion!

I'm looking to use Dion3 / NorDion2 for training hybrid architectures. While reviewing the codebase and documentation, I noticed a discrepancy regarding support for 3D+ tensors:

1. README: Mentions that tensors with 3+ dimensions (e.g., convolution layers) are currently not supported for Dion, with experimental flattening only available in Muon.
2. Code (dion/nordion2.py): NorDion2.__init__ explicitly exposes the flatten argument with the docstring:
   https://github.com/microsoft/dion/blob/main/dion/nordion2.py#L48

flatten: Whether to flatten 3D+ tensors to 2D for Muon updates.
True: Tensors with 3+ dimensions are flattened to 2D. Use this for convolutional layers.

Clarifying Questions:

• Is the README simply out of date?
• Or was the flatten parameter carried over from base/Muon implementations and not yet supported under Dion3's row selection / megabatching pipeline?

Context & Potential Contribution:

In modern RNN-based and hybrid architectures (e.g., in Mamba/GDN), most layers replace standard attention with recurrent or local 1D convolutional blocks. Therefore, without support for flattening 3D+ tensors into 2D matrices, a significant portion of the model's parameters cannot be updated with Dion/Muon. This would forces a messy split-optimizer setup (falling back to AdamW) and leaves a big fraction of the model's parameters on the table.

If flatten=True is not yet supported and you are open to community contributions, I'd be happy to explore putting together a PR to support flattening 3D+ tensors within the Dion pipeline.

Thanks!