Commit Graph

4 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
15a5e74434 Split VarLenTensor / SparseTensor to match Pixal3D's hierarchy
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.
2026-08-02 10:35:22 +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