Fix hlo_diff_test to compile the real NNX path

hlo_diff_test loaded each model yml directly as the top-level config, which
skips base.yml and leaves pure_nnx at its pydantic schema default (False), so
the test kept silently compiling the Linen path even after the NNX defaults
were flipped to True in base.yml. Its reference HLOs are therefore still Linen.

Compile via base.yml + model_name (the normal training path) so the config
inherits base.yml's logical_axis_rules and exercises the NNX path. The three
reference HLOs are regenerated for NNX via the Update HLO References workflow.

Also flip the enable_nnx/pure_nnx/pure_nnx_decoder pydantic schema defaults to
True so a config built without base.yml no longer silently falls back to the
Linen path.

Author: ecnal-cienet

src/maxtext/configs/types.py

@@ -988,11 +988,11 @@ class HardwareAndMesh(BaseModel):
       CustomRule.DEFAULT, description="Customized mesh and logical rules for granularity."
   )
   allow_split_physical_axes: bool = Field(False, description="Allow splitting physical axes for device mesh creation.")
-  enable_nnx: bool = Field(False, description="Whether to use NNX for model definition.")
+  enable_nnx: bool = Field(True, description="Whether to use NNX for model definition.")
   optimize_mesh_for_tpu_v6e: bool = Field(False, description="Apply transformations to the mesh for TPU v6e.")
   shardy: bool = Field(True, description="Whether to use shardy XLA backend.")
-  pure_nnx_decoder: bool = Field(False, description="Whether to enable pure NNX decoder.")
-  pure_nnx: bool = Field(False, description="Whether to enable pure NNX mode.")
+  pure_nnx_decoder: bool = Field(True, description="Whether to enable pure NNX decoder.")
+  pure_nnx: bool = Field(True, description="Whether to enable pure NNX mode.")
   remove_size_one_mesh_axis_from_type: bool = Field(
       True, description="Whether to remove size one mesh axis from type through jax.config."
   )

tests/integration/hlo_diff_test.py

@@ -93,11 +93,11 @@ def get_filtered_lines(file_path):
     return False
 
   @pytest.mark.parametrize(
-      "test_id, config_file, overrides",
+      "test_id, model_name, overrides",
       [
           (
               "deepseek3",
-              "src/maxtext/configs/models/deepseek3-test.yml",
+              "deepseek3-test",
               {
                   "compile_topology": "v6e-4",
                   "base_num_decoder_layers": 4,
@@ -107,7 +107,7 @@ def get_filtered_lines(file_path):
           ),
           (
               "llama3_8b",
-              "src/maxtext/configs/models/llama3-8b.yml",
+              "llama3-8b",
               {
                   "compile_topology": "v6e-4",
                   "base_num_decoder_layers": 4,
@@ -117,7 +117,7 @@ def get_filtered_lines(file_path):
           ),
           (
               "qwen3_1.7b",
-              "src/maxtext/configs/models/qwen3-1.7b.yml",
+              "qwen3-1.7b",
               {
                   "compile_topology": "v6e-4",
                   "base_num_decoder_layers": 4,
@@ -127,7 +127,7 @@ def get_filtered_lines(file_path):
           ),
       ],
   )
-  def test_hlo_diff(self, test_id, config_file, overrides):
+  def test_hlo_diff(self, test_id, model_name, overrides):
     """Test HLO diff for parameterized configurations."""
     local_landing_dir = os.path.join(os.path.dirname(__file__), f"hlo_diff_dump_{test_id}")
 
@@ -138,12 +138,16 @@ def test_hlo_diff(self, test_id, config_file, overrides):
 
     try:
       base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
-      config_path = os.path.join(base_dir, config_file)
+      # Load via base.yml + model_name, the normal training path: pyconfig resolves the model
+      # yml from model_name and merges it onto base.yml. Loading the model yml directly as the
+      # top-level config skips base.yml, leaving logical_axis_rules empty.
+      base_config_path = os.path.join(base_dir, "src/maxtext/configs/base.yml")
 
       # Arguments for train_compile
       test_args = [
           None,
-          config_path,
+          base_config_path,
+          f"model_name={model_name}",
           "dataset_type=synthetic",
           "override_model_config=true",
           "compile_topology_num_slices=1",