Add remove size one func to eval dataloader

[NuojCheng]

Description

Since remove_size_one_mesh_axis_from_type is tuned as False in batch split implementation, some redundant physical axes, e.g. context, data, are included in eval data iterators. This PR additionally enforces size one mesh axes get ignored when the eval data get generated.

Tests

Repro in v5p-8 vm: https://paste.googleplex.com/5394820295163904

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[RissyRan]

Thanks for the fix! Please ensure tests are passing as well.