Commit Graph

3 Commits

Author SHA1 Message Date
John
5dd3f11d92 Decoder ops: spatial<->channel, ConvNeXt block, C2S upsampling block
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.
2026-08-02 13:24:59 +10:00
John
cee423501c Add masked upsample and parameterise downsample reduction (16/16 tests)
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.
2026-08-02 10:42:20 +10:00
John
70cd436eb9 Shared TRELLIS-lineage sparse core in MLX
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.
2026-08-02 10:31:07 +10:00