Add default rule constants to types.py and add base.yml parity test

[bvandermoon]

Description

Currently, MaxText uses a dual-configuration system where base.yml supplies baseline dictionary keys and types.py defines the Pydantic schema. This causes duplicate maintenance whenever a config flag is added or modified.

This PR is Phase 1 of consolidating configuration management to use types.py as the single source of truth. This is making base.yml and types.py fully consistent with each other.

Key Changes

• Default Rule Constants in Python: Moved DEFAULT_MESH_AXES, DEFAULT_DATA_SHARDING, and the complete list of 60+ DEFAULT_LOGICAL_AXIS_RULES from base.yml into src/maxtext/configs/types.py.
• Pydantic Default Factories: Updated HardwareAndMesh and LayoutAndSharding to populate their fields using default_factory=lambda: copy.deepcopy(...).
• Exhaustive Parity Test: Added test_base_yml_types_parity in tests/unit/pyconfig_test.py to assert that:
  • 100% of keys in base.yml exist in types.MaxTextConfig.
  • Direct instantiation of types.MaxTextConfig() matches all default values from base.yml.

Tests

• python3 -m unittest tests/unit/pyconfig_test.py -k test_base_yml_types_parity
• pytest tests/unit/configs_test.py
• pytest tests/unit/configs_value_test.py

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.

[gemini-code-assist]

Code Review

This pull request updates default configuration values and structures in types.py, including extracting default mesh axes, data sharding, and logical axis rules into constants, and updates related unit tests. Feedback includes moving an inline import to the top of the file to comply with PEP 8 and removing a redundant duplicate call to os.environ.pop.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[NuojCheng]

Nice fix! Thank you Branden