Cross self attention switch

[entrpn]

• Overrides custom flash block sizes for cross attention to prevent extra padding on small cross attention sequence lengths.
• Same as [WAN] Use different sharding strategy for self and cross attention. #250 but setting the sharding at the NNX module level instead of the splash attention kernel function.

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[entrpn]

Just realized that I also need to support ring with these new axis rules. If attention=ring, then force logical axis rules to change in pyconfig such that ring attention is supported.