Fix shape mismatch and add DeepSeek4HyperHead in mhc.py

[olufiyin19]

Fix shape mismatch and add DeepSeek4HyperHead in mhc.py

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request updates the Dockerfile to install the Google Cloud CLI, increases retry limits and sleep duration in the ADK agent, and refactors the DeepSeek4 HyperHead implementation in mhc.py to use a normalized sigmoid-scaled mixing mechanism. The feedback suggests using a random variance scaling initializer for hc_fn to allow symmetry breaking, updating the ADK agent's docstring to match the new sleep duration, and using secure HTTPS and robust curl flags in the Dockerfile.

[gemini-code-assist]

The weight matrix hc_fn is initialized using default_scalar_init, which is a constant initializer (constant(0.01)). Using a constant initializer for a projection/weight matrix prevents symmetry breaking during training, which can severely limit representation learning and model capacity. It should be initialized using a random variance scaling initializer like nd_dense_init (similar to how weight matrices are initialized in ManifoldConstrainedHyperConnections).

    scale_init = nd_dense_init(1.0, "fan_in", "normal")
    self.hc_fn = nnx.Param(
        scale_init(
            self.rngs.params(),
            (self.hc_mult, self.hc_mult * config.emb_dim),
            self.weight_dtype,
            in_axis=1,
            out_axis=0,
        ),
        out_sharding=(None, None),
    )

[gemini-code-assist]

The default value of sleep_seconds was updated from 30 to 60, but the docstring still mentions "a 30-second sleep". Please update the docstring to reflect the new default value or make it generic.

Suggested change

	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."""
	def _send_message_with_retry(chat, prompt, max_retries=5, sleep_seconds=60):
	"""Sends a message to Gemini with retry and a sleep on 429 rate-limit/quota errors."""

[gemini-code-assist]

Using http instead of https for package repositories can expose the system to man-in-the-middle (MITM) attacks or interception. Additionally, running curl without -fsSL can fail silently or write error pages to the keyring file, which makes debugging harder. It is safer to use https and add -fsSL to the curl command.

    echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
    curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \