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/mhc.py

@@ -313,4 +313,22 @@ def __call__(
     return res_out + post_out, metadata
 
 
+from maxtext.layers import linears
 
+
+class DeepSeek4HyperHead(linears.DenseGeneral):
+  """DeepSeek V4 HyperHead for projecting expanded hidden states."""
+
+  def __init__(self, config: Config, mesh: Mesh, rngs: nnx.Rngs):
+    super().__init__(
+        in_features_shape=config.mhc_expansion_rate * config.emb_dim,
+        out_features_shape=config.emb_dim,
+        weight_dtype=config.weight_dtype,
+        kernel_axes=("mlp", "embed"),
+        rngs=rngs,
+    )
+
+  def __call__(self, x):
+    b, l, k, d = x.shape
+    x = jnp.reshape(x, (b, l, k * d))
+    return super().__call__(x)

src/maxtext/layers/nnx_decoders.py

@@ -1969,12 +1969,6 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in):
 
           if deepstack_visual_embeds is not None and lyr < len(deepstack_visual_embeds):
             visual_embeds = deepstack_visual_embeds[lyr]
-            if bidirectional_mask is not None and visual_embeds is not None:
-              y = deepstack_process(y, bidirectional_mask, visual_embeds)
-
-    assert isinstance(y, jax.Array)
-
-    # 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)
@@ -1983,6 +1977,14 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in):
         hidden_state = mhc_reduce(y)
     else:
       hidden_state = y
+    assert isinstance(y, jax.Array)
+
+    # After the final transformer layer, `y` holds the raw, un-normalized hidden state.
+    if getattr(cfg, "mhc_expansion_rate", 1) > 1:
+      # (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"):