Fix DeepSeek4HyperHead instantiation and usage in NNXDecoder

src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/Dockerfile

@@ -7,7 +7,9 @@ RUN apt-get update && apt-get install -y git curl ca-certificates gpg && \
     curl -fsSL https://cli.github.com/packages/githubcli-archive-keyring.gpg | dd of=/usr/share/keyrings/githubcli-archive-keyring.gpg && \
     chmod go+r /usr/share/keyrings/githubcli-archive-keyring.gpg && \
     echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | tee /etc/apt/sources.list.d/github-cli.list > /dev/null && \
-    apt-get update && apt-get install -y gh && \
+    echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
+    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \
+    apt-get update && apt-get install -y gh google-cloud-cli && \
     rm -rf /var/lib/apt/lists/*
 
 # Copy only requirements first to leverage Docker cache for heavy installations

src/maxtext/experimental/agent/ckpt_validation_pipeline/agent_sidecar/adk_agent.py

@@ -11,7 +11,7 @@
 logger = logging.getLogger(__name__)
 
 
-def _send_message_with_retry(chat, prompt, max_retries=3, sleep_seconds=30):
+def _send_message_with_retry(chat, prompt, max_retries=5, sleep_seconds=60):
   """Sends a message to Gemini with retry and a 30-second sleep on 429 rate-limit/quota errors."""
   for attempt in range(1, max_retries + 1):
     try:

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

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,13 +1967,11 @@ 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