mHC Operators

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

Manifold-Constrained Hyper-Connections

tileops.sequence_modeling.MHCPreFwdOp

The pre-layer half of Manifold-Constrained Hyper-Connections (mHC).

Reduces the width-expanded stream to the single tensor a layer consumes, and returns the residual that MHCPostFwdOp mixes the layer output back into. The expansion width is read from the shape of phi rather than passed in.

Layout: BSHD

__init__

__init__(
    alpha_pre,
    alpha_post,
    alpha_res,
    sinkhorn_repeat,
    sinkhorn_eps=0.02,
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• alpha_pre (float) –

  Manifest params.alpha_pre, the pre-layer mixing weight.

• alpha_post (float) –

  Manifest params.alpha_post, the post-layer mixing weight.

• alpha_res (float) –

  Manifest params.alpha_res, the residual mixing weight.

• sinkhorn_repeat (int) –

  Manifest params.sinkhorn_repeat, Sinkhorn iterations.

• sinkhorn_eps (float, default: 0.02 ) –

  Manifest params.sinkhorn_eps, Sinkhorn entropy scale.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    phi,
    x,
    b,
)

Run the op on the inputs the manifest declares.

Parameters:

• phi (Tensor) –

  [n * c_x, n * n + 2 * n], dtype float32.

• x (Tensor) –

  [batch, n * c_x], dtype bfloat16.

• b (Tensor) –

  [n * n + 2 * n], dtype float32.

Returns:

• Tensor –

  x_res, x_layer, h_post, as the manifest declares.

tileops.sequence_modeling.MHCPostFwdOp

The post-layer half of Manifold-Constrained Hyper-Connections (mHC).

Mixes a layer's output into the residual MHCPreFwdOp set aside, restoring the width-expanded stream.

Layout: BSHD

__init__

__init__(
    *,
    target=None
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

forward

forward(
    x_layer_out,
    h_post,
    x_res,
)

Run the op on the inputs the manifest declares.

Parameters:

• x_layer_out (Tensor) –

  [batch, c_x], dtype bfloat16.

• h_post (Tensor) –

  [batch, n], dtype float32.

• x_res (Tensor) –

  [batch, n * c_x], dtype bfloat16.

Returns:

• Tensor –

  x_out, as the manifest declares.