diff --git a/src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/mock_failure_log.txt b/src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/mock_failure_log.txt
new file mode 100644
index 0000000000..da609f6d82
--- /dev/null
+++ b/src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/mock_failure_log.txt
@@ -0,0 +1,7 @@
+Traceback (most recent call last):
+  File "train.py", line 42, in <module>
+    import maxtext
+  File "/usr/local/google/home/fiyinbenstowe/Desktop/Project/maxtext/src/maxtext/layers/normalizations.py", line 72
+    mean2 = jnp.mean(lax.square(x), axis=-1, keepdims=True)
+          ^
+SyntaxError: invalid syntax
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 895ea27c14..e6c0fc6f8a 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -438,13 +438,6 @@ def __init__(
     self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4
     self.is_gemma4_small = self.config.decoder_block == DecoderBlockType.GEMMA4_SMALL
 
-    if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4:
-      self.hc_head = mhc.DeepSeek4HyperHead(
-          config=config,
-          mesh=self.mesh,
-          rngs=self.rngs,
-      )
-
     self._init_decoder_layers(decoder_block_classes, rngs, mesh)
 
   def _init_decoder_layers(self, decoder_block_classes, rngs, mesh):
@@ -1974,16 +1967,13 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in):
 
     assert isinstance(y, jax.Array)
 
+    # After the final transformer layer, `y` holds the raw, un-normalized hidden state.
     # After the final transformer layer, `y` holds the raw, un-normalized hidden state.
     if getattr(cfg, "mhc_expansion_rate", 1) > 1:
-      if cfg.decoder_block == DecoderBlockType.DEEPSEEK4:
-        hidden_state = self.hc_head(y)
-      else:
-        # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim)
-        hidden_state = mhc_reduce(y)
+      # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim)
+      hidden_state = mhc_reduce(y)
     else:
       hidden_state = y
-
     # When invoking from vLLM with RPA attention, logit computation is deferred to a later stage.
     if cfg.attention in ("vllm_rpa", "vllm_batched_rpa"):
       logits = None