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.
Pixal3D refactored the container into a VarLenTensor base (feats + explicit
slice-per-batch layout, no coords) with SparseTensor(VarLenTensor) adding
coordinates. LATO.2 has only the combined class. Modelling the split here means one
package serves both without either port adapting at every call site.
Also exposes Pixal3D's cache spelling (get/register_spatial_cache) alongside LATO.2's
(cache_get/cache_put). 13/13 tests unchanged.
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.
