Completes the sparse op surface Pixal3D's decoders need.
spatial2channel/channel2spatial are sparse space-to-depth and its inverse. The slot
index is sum_i (coord[i] % f) * f**i - axis 0 FASTEST varying, the reverse of C order.
decodes it the same way; using C order in one and this in the other would
misplace every child while keeping all shapes valid.
SparseResBlockC2S3d widens to out_channels*8 so channel2spatial can redistribute those
channels across the 8 children, with a predicted subdiv mask deciding which children
exist - so the occupied set grows selectively (32 voxels -> 122, not 256).
Tests cover exact round-trip, zero-fill of unoccupied slots, and selective growth.
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.
Two silent bugs, both found by diffing against upstream running on CPU torch. The flow
models have no sparse conv, so upstream is runnable here with flash-attn swapped for
SDPA - a real numerical oracle, unlike the sparse path.
1. rope_phases ships as COMPLEX64 (torch.polar), so the rotation is a complex multiply
and cos/sin are the phase's real/imag parts. Taking cos() of a complex phase was
completely wrong. Verified against torch's view_as_complex formulation to 1.2e-7.
2. Upstream applies qk RMS norm BEFORE RoPE; I had it reversed. They do not commute -
RMS applies a per-component gain, RoPE rotates within pairs. Reversed, one block
still correlated 0.9998, which compounded to 0.84 across 30 blocks.
After both: SparseStructureFlowModel matches upstream at correlation 1.00000000,
max abs diff 1.2e-5, on the real 1.3B checkpoint.
Unlocks all four Pixal3D flow checkpoints (~20GB) at once - they are pure transformers
with no sparse conv. LATO.2 has flow models too (vertex_structured_flow, topo_flow), so
these belong in the shared core rather than either port.
Three details taken from upstream rather than assumed, each silent when wrong:
- norm1/norm3 are NON-affine but norm2 IS affine in the modulated cross block. There is
an explicit test asserting that asymmetry.
- MultiHeadRMSNorm is written upstream as F.normalize(x)*gamma*sqrt(dim). F.normalize is
L2, and the sqrt(d) turns it into RMS - implemented directly as RMS and verified equal
to the upstream formulation to 9.5e-7.
- RoPE phases are NOT derived: Pixal3D ships rope_phases as a stored tensor, so they are
passed in. Tested that the rotation preserves per-pair norms and is not a no-op.
Also tested: gates at zero make the block an identity on its residual branches.
Closes the op gap for Pixal3D.
downsample now takes mode='max'|'mean'. This differs BETWEEN MODELS and is silent when
wrong: LATO.2 hardcodes amax, Pixal3D parameterises it and defaults to mean. Callers
must be explicit.
upsample is a MASKED expansion, distinct from subdivide: subdivision.feats is an
[N, factor**3] 0/1 mask, so a voxel emits 0-8 children rather than always 8. Pixal3D's
decoder uses it to grow the occupied set selectively. Child-index decoding is
subidx // factor**i % factor - axis 0 fastest-varying, the REVERSE of subdivide's
C-order, which is exactly the kind of mixup that misplaces every child silently.
Tests cover both modes differing, masked expansion landing at the right offsets, and a
voxel with no flagged children contributing nothing.
Extracted from lato.2_mrp_mlx. The same sparse module underlies LATO.2, Pixal3D and
the rest of the TRELLIS.2 family, and all of them are blocked on Apple Silicon by the
same single op - submanifold conv - so it belongs in one tested package rather than
vendored per port.
13/13 tests: 7 for the conv against a hand-written reference (spconv is uninstallable
here so there is no upstream oracle), 6 for the remaining layers against torch.
SubMConv3d runs 13.8ms at 128^3/128ch on m3ultra, 5.9x faster than the obvious
per-offset loop.
