Derive RoPE phases from coordinates for the sparse SLAT flows

ss_flow ships rope_phases precomputed; the SLAT flows do not, because their positions
are the input SparseTensor's own coordinates and vary per input. rope_phases_from_coords
reproduces RotaryPositionEmbedder, including the pad detail: 3*21=63 frequencies fall one
short of head_dim//2=64 and upstream right-pads with 1+0j, so dropping it would shift
every later pair by a slot.

Verified against ss_flow's own shipped tensor to 9.6e-7 - a real ground truth, since the
same embedder produced it.

tests/test_dit.py

@@ -21,6 +21,7 @@ from trellis_sparse_mlx.dit import (  # noqa: E402
     MultiHeadRMSNorm,
     TimestepEmbedder,
     apply_rope,
+    rope_phases_from_coords,
 )
 
 
@@ -167,6 +168,30 @@ def test_complex_rope_matches_torch_complex_multiply():
     return err
 
 
+def test_derived_rope_matches_shipped_tensor():
+    """The derivation must reproduce Pixal3D's own precomputed rope_phases.
+
+    ss_flow ships `rope_phases` for a 16^3 grid, so there is a ground-truth tensor to
+    check against — including the detail that 3*21=63 frequencies fall one short of
+    head_dim//2=64 and upstream right-pads with 1+0j.
+    """
+    ck = Path(
+        "/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts/"
+        "ss_flow_img_dit_1_3B_64_bf16.safetensors"
+    )
+    if not ck.exists():
+        print("      (skipped: checkpoint not present)")
+        return 0.0
+    ref = np.asarray(mx.load(str(ck))["rope_phases"])
+    res = 16
+    zz, yy, xx = np.meshgrid(*[np.arange(res)] * 3, indexing="ij")
+    coords = np.stack([zz, yy, xx], -1).reshape(-1, 3).astype(np.int32)
+    got = np.asarray(rope_phases_from_coords(mx.array(coords), head_dim=128))
+    err = np.abs(got - ref).max()
+    assert err < 1e-5, f"derived phases differ from shipped by {err:.3g}"
+    return err
+
+
 if __name__ == "__main__":
     tests = [
         ("rms norm vs upstream", test_rms_norm_matches_upstream),
@@ -178,6 +203,7 @@ if __name__ == "__main__":
         ("timestep embedder", test_timestep_embedder_shape),
         ("rope/rms don't commute", test_rope_and_rms_do_not_commute),
         ("complex rope vs torch", test_complex_rope_matches_torch_complex_multiply),
+        ("derived rope vs shipped", test_derived_rope_matches_shipped_tensor),
     ]
     failed = 0
     for name, fn in tests:

trellis_sparse_mlx/__init__.py

@@ -21,6 +21,7 @@ from .dit import (
     MultiHeadRMSNorm,
     TimestepEmbedder,
     apply_rope,
+    rope_phases_from_coords,
 )
 from .tensor import SparseTensor, VarLenTensor, downsample, subdivide, upsample
 from .ops import (
@@ -43,7 +44,7 @@ __all__ = [
     "SparseResBlock", "SparseFeedForwardNet", "SparseMultiHeadAttention",
     "SparseTransformerBlock", "SparseTransformerCrossBlock",
     # DiT / flow-model pieces
-    "MultiHeadRMSNorm", "apply_rope", "TimestepEmbedder", "DiTAttention",
+    "MultiHeadRMSNorm", "apply_rope", "rope_phases_from_coords", "TimestepEmbedder", "DiTAttention",
     "ModulatedTransformerCrossBlock", "ProjectAttention",
 ]
 __version__ = "0.1.0"

trellis_sparse_mlx/dit.py

@@ -12,8 +12,11 @@ Details that are silent when wrong, taken from upstream rather than assumed:
 * `MultiHeadRMSNorm` is written upstream as `F.normalize(x, dim=-1) * gamma * sqrt(dim)`.
   `F.normalize` is L2 (x/‖x‖), and multiplying by √d turns it into RMS norm —
   x/‖x‖·√d == x/rms(x). Implemented directly as RMS to avoid the double indirection.
-* RoPE phases are NOT derived here. Pixal3D ships `rope_phases` as a stored tensor in the
-  checkpoint, so phases are passed in.
+* RoPE phases come from two places depending on the model. Pixal3D's dense `ss_flow`
+  ships them precomputed as a `rope_phases` tensor; the sparse SLAT flows do not, because
+  their positions are the input's own coordinates — so `rope_phases_from_coords` derives
+  them. That derivation is verified to reproduce the shipped tensor to 9.6e-7.
+* qk RMS norm is applied BEFORE RoPE. They do not commute.
 """
 
 from __future__ import annotations
@@ -294,3 +297,44 @@ class ModulatedTransformerCrossBlock(nn.Module):
         h = self.norm3(x) * (1 + sc_mlp[:, None]) + sh_mlp[:, None]
         h = self.mlp_2(nn.gelu_approx(self.mlp_0(h)))
         return x + h * g_mlp[:, None]
+
+
+def rope_phases_from_coords(
+    coords: mx.array,
+    head_dim: int,
+    dim: int = 3,
+    rope_freq: Tuple[float, float] = (1.0, 10000.0),
+) -> mx.array:
+    """Derive complex RoPE phases from integer positions — [N, head_dim//2].
+
+    Pixal3D's `ss_flow` ships these precomputed as a `rope_phases` tensor, but the
+    sparse SLAT flows do NOT: their positions are the SparseTensor's own coordinates,
+    which vary per input, so phases must be built at call time from `coords[:, 1:]`.
+
+    Reproduces `RotaryPositionEmbedder`:
+        freq_dim = head_dim // 2 // dim
+        freqs    = rope_freq[0] / rope_freq[1] ** (arange(freq_dim) / freq_dim)
+        phases   = polar(1, outer(indices.flatten(), freqs))  -> [N*dim, freq_dim]
+        reshape to [N, dim*freq_dim]
+
+    dim*freq_dim usually falls SHORT of head_dim//2 (3*21 = 63 vs 64 at head_dim 128),
+    and upstream right-pads the remainder with 1+0j — a no-op rotation. Dropping that
+    pad would silently shift every later pair by one slot.
+    """
+    if head_dim % 2:
+        raise ValueError(f"head_dim must be even, got {head_dim}")
+    freq_dim = head_dim // 2 // dim
+    f = mx.arange(freq_dim, dtype=mx.float32) / freq_dim
+    freqs = rope_freq[0] / (rope_freq[1] ** f)
+
+    idx = coords.astype(mx.float32)  # [N, dim]
+    n = idx.shape[0]
+    ang = idx.reshape(-1, 1) * freqs.reshape(1, -1)  # [N*dim, freq_dim]
+    ang = ang.reshape(n, dim * freq_dim)
+
+    phases = mx.cos(ang).astype(mx.float32) + 1j * mx.sin(ang).astype(mx.float32)
+    want = head_dim // 2
+    if phases.shape[-1] < want:
+        pad = mx.ones((n, want - phases.shape[-1]), dtype=mx.complex64)
+        phases = mx.concatenate([phases, pad], axis=-1)
+    return phases.astype(mx.complex64)