Fix rank mismatch in MaxText synthetic data sharding

[lukebaumann]

Description

This PR fixes a rank mismatch issue in MaxText synthetic data sharding during data loading.

Root Cause

SyntheticDataIterator was using the legacy config.data_sharding which resolved to a 1D sharding spec P(('data', 'fsdp')) (after filtering). When applied to 2D output tensors of shape (batch, seq), JAX sharding validation failed with AssertionError: (1, 2) (rank mismatch) on JAX builds that strictly enforce this check.

Solution

Modified SyntheticDataIterator to use sharding.get_input_data_sharding(config, mesh). This helper uses config.input_data_sharding_logical_axes which correctly resolves to a 2D sharding spec P(('data', 'fsdp'), None), matching the rank of the output tensors.

Also removed the unused PartitionSpec as P import in synthetic_data_processing.py.

Tests

Added a new unit test tests/unit/synthetic_data_test.py which is parameterized to test both auto and explicit shard_mode:

1. Forces 4 CPU devices.
2. Creates a 2x2 mesh.
3. Initializes SyntheticDataIterator with llama3.1-8b config.
4. Verifies the output shape is (8, 16) and sharding is exactly P(('data', 'fsdp'), None).

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!