For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

HyperConnection

HyperConnection

class max.experimental.nn.HyperConnection(hidden_size, hc_mult, *, hc_sinkhorn_iters=20, hc_eps=1e-06, rms_norm_eps=1e-06)

source

Bases: Module

Mixes hc_mult parallel residual streams around a transformer sublayer.

Manifold-Constrained Hyper-Connections (mHC) generalize the residual connection: instead of one stream carrying the residual between sublayers, hc_mult streams run in parallel, and a learned gate decides how they collapse into the sublayer’s input and how the sublayer’s output is written back. Reference: Xie et al. 2026, section 2.2 equation 8.

A single projection of the normalized streams produces the three gates:

• pre: how the parallel streams collapse into the single sequence the sublayer consumes. Applied here, producing collapsed.
• post: where the sublayer’s output lands across the streams, in [0, 2]. Returned for the caller to apply.
• comb: an hc_mult x hc_mult mixer over the streams, projected onto the doubly-stochastic manifold by Sinkhorn-Knopp. Returned for the caller to apply.

A decoder layer holds two of these, one for the attention site and one for the feed-forward site. The caller drives the surrounding residual update:

post, comb, collapsed = hc(hidden_streams)
sublayer_out = sublayer(collapsed)
hidden_streams = comb @ hidden_streams + post.unsqueeze(-1) * sublayer_out

The gate math always runs in float32. collapsed is cast back to the input’s dtype.

Parameters:

• hidden_size (int) – Width of a single residual stream.
• hc_mult (int) – Number of parallel residual streams. Must be a power of two whose square fits in one warp.
• hc_sinkhorn_iters (int) – Sinkhorn-Knopp iterations used to project comb.
• hc_eps (float) – Epsilon guarding the Sinkhorn divisions.
• rms_norm_eps (float) – Epsilon of the unweighted RMSNorm applied to the streams before the projection.

forward()

forward(hidden_streams)

source

Computes the mHC gates and collapses the streams for the sublayer.

Parameters:

hidden_streams (Tensor) – The parallel residual streams, shaped [total_seq_len, hc_mult, hidden_size]. The batch and sequence axes are folded into the leading axis.

Returns:

A tuple (post, comb, collapsed). post is float32 [total_seq_len, hc_mult], comb is float32 [total_seq_len, hc_mult, hc_mult], and collapsed is [total_seq_len, hidden_size] in hidden_streams’ dtype.

Raises:

ValueError – If hidden_streams is not rank 3 or its trailing two dimensions do not match hc_mult and hidden_size.

Return type:

tuple[Tensor, Tensor, Tensor]

hc_base

hc_base: Tensor

source

Per-output bias of shape [2 * hc_mult + hc_mult ** 2], in pre, post, comb order.

hc_fn

hc_fn: Tensor

source

Stream projection of shape [2 * hc_mult + hc_mult ** 2, hc_mult * hidden_size].

hc_scale

hc_scale: Tensor

source

Per-output scale of shape [3], in pre, post, comb order.