Merge pull request #4378 from CIeNET-International:emma/fix-gemma4-cp-convert

PiperOrigin-RevId: 944551277

Author: Google-ML-Automation

src/maxtext/layers/nnx_decoders.py

@@ -653,7 +653,7 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh):
     RemattedGemma4Block = gemma4.Gemma4ScannableBlock
 
     if scan_length > 0:
-      self.layers = self._create_scanned_layers(
+      self.scanned_blocks = self._create_scanned_layers(
           RemattedGemma4Block,
           length=scan_length,
           metadata_axis_name="layers",
@@ -2030,8 +2030,8 @@ def _apply_gemma4_scanned_blocks(
       grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan(
           kv_caches, scan_length, attention_pattern_length, stack=False
       )
-      y, self.layers, _ = self._apply_layers_sequentially(
-          self.layers, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, **layer_kwargs
+      y, self.scanned_blocks, _ = self._apply_layers_sequentially(
+          self.scanned_blocks, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, **layer_kwargs
       )
       maxtext_utils.update_kv_caches_after_scan(
           kv_caches, grouped_kv_caches, scan_length, attention_pattern_length, stacked=False