Fix: Numeric divergence between the mHC layer and its Pallas kernel

[denis-mil]

Bug report: #5195

Description

Motivation:
The reference path had its own mapping helper that cast alpha_scale and
beta down to self.dtype before computing sigmoid(alpha_scale * h + beta),
while the Pallas path called common.compute_sigmoid_gate, which keeps the
scale and bias in float32. With dtype: bfloat16 the two branches therefore
gated on different numbers, so toggling use_mhc_pallas_kernel changed the
layer's output beyond accumulation noise.

Implementation:
Drop ManifoldConstrainedHyperConnections.mapping and call the shared
compute_sigmoid_gate from both pre- and post-mapping, casting the float32
result back to self.dtype so the following GEMM is unchanged. Export the
helper from maxtext.kernels.mhc so the layer can reach it.

Tests:

• test_pallas_kernel_matches_reference_path runs the same weights through
  both branches and asserts they agree; the _tpu variant repeats it against
  the Mosaic-compiled kernel instead of interpret mode.
• test_sigmoid_gate_computes_in_float32 pins the helper's precision directly,
  since the end-to-end comparison cannot resolve a bfloat16 regression there.
• Fold the four copies of the interpret-mode kernel patching into the
  _interpreted_mhc_kernel context manager, and the repeated module/MLP
  construction into _build_mhc_and_mlp.

Tests

Unit tested on CPU and TPU v6e-1

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request refactors the Manifold-Constrained Hyper Connections (mHC) implementation by sharing the compute_sigmoid_gate helper function between the layer and the Pallas kernel, ensuring identical gating behavior across both paths. The local mapping method in src/maxtext/layers/mhc.py has been removed in favor of this shared helper, which computes in float32 and is subsequently cast back to the activation dtype. Additionally, the unit tests in tests/unit/mhc_test.py have been refactored to reduce duplication using helper methods and a new context manager for kernel interpretation, and new tests have been added to verify that the Pallas kernel matches the reference path and that the sigmoid gate computes correctly in float32. There are no review comments, so no further feedback is provided.