5dd3f11d92
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.
332 lines
13 KiB
Python
332 lines
13 KiB
Python
"""The rest of the sparse layer set, in MLX.
|
|
|
|
Nothing here is CUDA-locked upstream — these are ordinary linear/norm/attention layers
|
|
that merely take a SparseTensor instead of a dense one. They are reimplemented rather
|
|
than adapted because upstream's versions inherit from torch modules.
|
|
|
|
Two normalisation shapes are easy to conflate, and upstream uses both:
|
|
|
|
LayerNorm32 applied to `x.feats` directly -> per-voxel over channels.
|
|
SparseGroupNorm32 reshapes [N_b, C] -> [1, C, N_b] per batch item, so statistics
|
|
are over (channels-in-group x voxels) WITHIN one batch item.
|
|
Getting this wrong is silent: shapes match either way.
|
|
|
|
Attention runs in `attn_mode="full"`, which upstream defines as full attention *within*
|
|
each batch item (never across). Batch rows are contiguous, so each item is one slice.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Optional
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
|
|
from .conv import SubMConv3d
|
|
from .tensor import SparseTensor
|
|
|
|
# ---------------------------------------------------------------- primitives
|
|
|
|
|
|
class SparseLinear(nn.Module):
|
|
def __init__(self, in_features: int, out_features: int, bias: bool = True):
|
|
super().__init__()
|
|
self.linear = nn.Linear(in_features, out_features, bias=bias)
|
|
|
|
def __call__(self, x: SparseTensor) -> SparseTensor:
|
|
return x.replace(self.linear(x.feats))
|
|
|
|
|
|
class LayerNorm32(nn.Module):
|
|
"""Per-voxel LayerNorm over channels; computed in fp32 as upstream does."""
|
|
|
|
def __init__(self, dim: int, affine: bool = False, eps: float = 1e-6):
|
|
super().__init__()
|
|
self.eps = eps
|
|
self.affine = affine
|
|
if affine:
|
|
self.weight = mx.ones((dim,))
|
|
self.bias = mx.zeros((dim,))
|
|
|
|
def __call__(self, feats: mx.array) -> mx.array:
|
|
dt = feats.dtype
|
|
f = feats.astype(mx.float32)
|
|
mu = mx.mean(f, axis=-1, keepdims=True)
|
|
var = mx.var(f, axis=-1, keepdims=True)
|
|
f = (f - mu) * mx.rsqrt(var + self.eps)
|
|
if self.affine:
|
|
f = f * self.weight + self.bias
|
|
return f.astype(dt)
|
|
|
|
|
|
class SparseGroupNorm32(nn.Module):
|
|
"""GroupNorm over (channels-in-group x voxels), per batch item. fp32 internally."""
|
|
|
|
def __init__(self, num_groups: int, num_channels: int, eps: float = 1e-5):
|
|
super().__init__()
|
|
if num_channels % num_groups != 0:
|
|
raise ValueError(f"{num_channels} channels not divisible by {num_groups}")
|
|
self.num_groups = num_groups
|
|
self.num_channels = num_channels
|
|
self.eps = eps
|
|
self.weight = mx.ones((num_channels,))
|
|
self.bias = mx.zeros((num_channels,))
|
|
|
|
def __call__(self, x: SparseTensor) -> SparseTensor:
|
|
dt = x.feats.dtype
|
|
g, c = self.num_groups, self.num_channels
|
|
parts = []
|
|
for sl in x.layout:
|
|
f = x.feats[sl].astype(mx.float32) # [n_b, C]
|
|
n_b = f.shape[0]
|
|
if n_b == 0:
|
|
parts.append(f)
|
|
continue
|
|
# -> [G, (C/G)*n_b] so mean/var cover channels *and* voxels in the group
|
|
grouped = f.T.reshape(g, (c // g) * n_b)
|
|
mu = mx.mean(grouped, axis=1, keepdims=True)
|
|
var = mx.var(grouped, axis=1, keepdims=True)
|
|
grouped = (grouped - mu) * mx.rsqrt(var + self.eps)
|
|
f = grouped.reshape(c, n_b).T
|
|
parts.append(f * self.weight + self.bias)
|
|
out = parts[0] if len(parts) == 1 else mx.concatenate(parts, axis=0)
|
|
return x.replace(out.astype(dt))
|
|
|
|
|
|
class SparseSiLU(nn.Module):
|
|
def __call__(self, x: SparseTensor) -> SparseTensor:
|
|
return x.replace(nn.silu(x.feats))
|
|
|
|
|
|
class SparseGELU(nn.Module):
|
|
def __call__(self, x: SparseTensor) -> SparseTensor:
|
|
return x.replace(nn.gelu(x.feats))
|
|
|
|
|
|
# ---------------------------------------------------------------- blocks
|
|
|
|
|
|
class SparseResBlock(nn.Module):
|
|
"""norm1(affine) -> silu -> conv1 -> norm2(no affine) -> silu -> conv2 + skip."""
|
|
|
|
def __init__(self, channels: int, out_channels: Optional[int] = None):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.out_channels = out_channels or channels
|
|
self.norm1 = LayerNorm32(channels, affine=True, eps=1e-6)
|
|
self.norm2 = LayerNorm32(self.out_channels, affine=False, eps=1e-6)
|
|
self.conv1 = SubMConv3d(channels, self.out_channels, 3)
|
|
self.conv2 = SubMConv3d(self.out_channels, self.out_channels, 3)
|
|
self.skip_connection = (
|
|
SparseLinear(channels, self.out_channels)
|
|
if channels != self.out_channels
|
|
else None
|
|
)
|
|
|
|
def __call__(self, x: SparseTensor) -> SparseTensor:
|
|
h = x.replace(self.norm1(x.feats))
|
|
h = h.replace(nn.silu(h.feats))
|
|
h = self.conv1(h)
|
|
h = h.replace(self.norm2(h.feats))
|
|
h = h.replace(nn.silu(h.feats))
|
|
h = self.conv2(h)
|
|
skip = self.skip_connection(x).feats if self.skip_connection else x.feats
|
|
return h.replace(h.feats + skip)
|
|
|
|
|
|
class SparseFeedForwardNet(nn.Module):
|
|
"""Upstream is nn.Sequential(Linear, GELU, Linear), so its checkpoint keys are
|
|
`mlp.mlp.0` and `mlp.mlp.2` — index 1 is the activation and carries no weights.
|
|
Named `mlp_0`/`mlp_2` here because a Python list with a None hole does not survive
|
|
MLX's parameter tree; the loader remaps the dotted indices onto these."""
|
|
|
|
def __init__(self, channels: int, mlp_ratio: float = 4.0):
|
|
super().__init__()
|
|
hidden = int(channels * mlp_ratio)
|
|
self.mlp_0 = nn.Linear(channels, hidden)
|
|
self.mlp_2 = nn.Linear(hidden, channels)
|
|
|
|
def __call__(self, x: SparseTensor) -> SparseTensor:
|
|
return x.replace(self.mlp_2(nn.gelu_approx(self.mlp_0(x.feats))))
|
|
|
|
|
|
def _sdpa_per_batch(
|
|
q: mx.array, k: mx.array, v: mx.array, layout_q, layout_kv, heads: int, scale: float
|
|
) -> mx.array:
|
|
"""Full attention inside each batch item. q/k/v are [N, H, D] flattened over batch."""
|
|
outs = []
|
|
for sq, skv in zip(layout_q, layout_kv):
|
|
qi = q[sq].transpose(1, 0, 2)[None] # [1, H, n, D]
|
|
ki = k[skv].transpose(1, 0, 2)[None]
|
|
vi = v[skv].transpose(1, 0, 2)[None]
|
|
o = mx.fast.scaled_dot_product_attention(qi, ki, vi, scale=scale)
|
|
outs.append(o[0].transpose(1, 0, 2)) # [n, H, D]
|
|
return outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=0)
|
|
|
|
|
|
class SparseMultiHeadAttention(nn.Module):
|
|
"""attn_mode='full' only — the sole mode LATO.2's model code instantiates."""
|
|
|
|
def __init__(
|
|
self,
|
|
channels: int,
|
|
num_heads: int,
|
|
ctx_channels: Optional[int] = None,
|
|
attn_type: str = "self",
|
|
qkv_bias: bool = True,
|
|
):
|
|
super().__init__()
|
|
if channels % num_heads != 0:
|
|
raise ValueError(f"{channels} channels not divisible by {num_heads} heads")
|
|
self.channels = channels
|
|
self.num_heads = num_heads
|
|
self.head_dim = channels // num_heads
|
|
self.scale = self.head_dim**-0.5
|
|
self._type = attn_type
|
|
self.ctx_channels = ctx_channels if ctx_channels is not None else channels
|
|
if attn_type == "self":
|
|
self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
|
|
else:
|
|
self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
|
|
self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
|
|
self.to_out = nn.Linear(channels, channels)
|
|
|
|
def __call__(
|
|
self, x: SparseTensor, context: Optional[SparseTensor] = None
|
|
) -> SparseTensor:
|
|
n = x.feats.shape[0]
|
|
h, d = self.num_heads, self.head_dim
|
|
if self._type == "self":
|
|
qkv = self.to_qkv(x.feats).reshape(n, 3, h, d)
|
|
q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]
|
|
lq = lkv = x.layout
|
|
else:
|
|
if context is None:
|
|
raise ValueError("cross-attention needs a context")
|
|
q = self.to_q(x.feats).reshape(n, h, d)
|
|
m = context.feats.shape[0]
|
|
kv = self.to_kv(context.feats).reshape(m, 2, h, d)
|
|
k, v = kv[:, 0], kv[:, 1]
|
|
lq, lkv = x.layout, context.layout
|
|
o = _sdpa_per_batch(q, k, v, lq, lkv, h, self.scale)
|
|
return x.replace(self.to_out(o.reshape(n, self.channels)))
|
|
|
|
|
|
class SparseTransformerBlock(nn.Module):
|
|
"""Pre-norm self-attention + FFN. Norms are non-affine (ln_affine=False upstream)."""
|
|
|
|
def __init__(self, channels: int, num_heads: int, mlp_ratio: float = 4.0):
|
|
super().__init__()
|
|
self.norm1 = LayerNorm32(channels, affine=False, eps=1e-6)
|
|
self.norm2 = LayerNorm32(channels, affine=False, eps=1e-6)
|
|
self.attn = SparseMultiHeadAttention(channels, num_heads)
|
|
self.mlp = SparseFeedForwardNet(channels, mlp_ratio)
|
|
|
|
def __call__(self, x: SparseTensor) -> SparseTensor:
|
|
h = self.attn(x.replace(self.norm1(x.feats)))
|
|
x = x.replace(x.feats + h.feats)
|
|
h = self.mlp(x.replace(self.norm2(x.feats)))
|
|
return x.replace(x.feats + h.feats)
|
|
|
|
|
|
class SparseTransformerCrossBlock(nn.Module):
|
|
"""Pre-norm self-attn -> cross-attn -> FFN."""
|
|
|
|
def __init__(
|
|
self, channels: int, ctx_channels: int, num_heads: int, mlp_ratio: float = 4.0
|
|
):
|
|
super().__init__()
|
|
self.norm1 = LayerNorm32(channels, affine=False, eps=1e-6)
|
|
self.norm2 = LayerNorm32(channels, affine=False, eps=1e-6)
|
|
self.norm3 = LayerNorm32(channels, affine=False, eps=1e-6)
|
|
self.context_norm = LayerNorm32(ctx_channels, affine=False, eps=1e-6)
|
|
self.self_attn = SparseMultiHeadAttention(channels, num_heads)
|
|
self.cross_attn = SparseMultiHeadAttention(
|
|
channels, num_heads, ctx_channels=ctx_channels, attn_type="cross"
|
|
)
|
|
self.mlp = SparseFeedForwardNet(channels, mlp_ratio)
|
|
|
|
def __call__(self, x: SparseTensor, context: SparseTensor) -> SparseTensor:
|
|
h = self.self_attn(x.replace(self.norm1(x.feats)))
|
|
x = x.replace(x.feats + h.feats)
|
|
ctx = context.replace(self.context_norm(context.feats))
|
|
h = self.cross_attn(x.replace(self.norm2(x.feats)), ctx)
|
|
x = x.replace(x.feats + h.feats)
|
|
h = self.mlp(x.replace(self.norm3(x.feats)))
|
|
return x.replace(x.feats + h.feats)
|
|
|
|
|
|
class SparseConvNeXtBlock3d(nn.Module):
|
|
"""conv -> norm -> MLP, residual. The decoder workhorse in Pixal3D.
|
|
|
|
Note the ordering: convolution comes FIRST, before the norm — unlike the pre-norm
|
|
SparseResBlock. And the residual adds the block input, not the post-conv tensor.
|
|
"""
|
|
|
|
def __init__(self, channels: int, mlp_ratio: float = 4.0):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.norm = LayerNorm32(channels, affine=True, eps=1e-6)
|
|
self.conv = SubMConv3d(channels, channels, 3)
|
|
hidden = int(channels * mlp_ratio)
|
|
self.mlp_0 = nn.Linear(channels, hidden)
|
|
self.mlp_2 = nn.Linear(hidden, channels)
|
|
|
|
def __call__(self, x: SparseTensor) -> SparseTensor:
|
|
h = self.conv(x)
|
|
h = h.replace(self.norm(h.feats))
|
|
h = h.replace(self.mlp_2(nn.silu(self.mlp_0(h.feats))))
|
|
return h.replace(h.feats + x.feats)
|
|
|
|
|
|
class SparseResBlockC2S3d(nn.Module):
|
|
"""Upsampling residual block: conv to 8x channels, then channel->spatial.
|
|
|
|
`conv1` widens to `out_channels * 8` so that `channel2spatial` can redistribute those
|
|
channels into the 8 children of each voxel. Which children exist is decided by the
|
|
`subdiv` mask — predicted by `to_subdiv` when `pred_subdiv` is set — so the occupied
|
|
set grows selectively rather than always x8.
|
|
|
|
The skip path is a `repeat_interleave`, not a linear: the input is upsampled by the
|
|
same channel->spatial step and its channels are repeated to match `out_channels`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
channels: int,
|
|
out_channels: Optional[int] = None,
|
|
pred_subdiv: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.out_channels = out_channels or channels
|
|
self.pred_subdiv = pred_subdiv
|
|
self.norm1 = LayerNorm32(channels, affine=True, eps=1e-6)
|
|
self.norm2 = LayerNorm32(self.out_channels, affine=False, eps=1e-6)
|
|
self.conv1 = SubMConv3d(channels, self.out_channels * 8, 3)
|
|
self.conv2 = SubMConv3d(self.out_channels, self.out_channels, 3)
|
|
if pred_subdiv:
|
|
self.to_subdiv = SparseLinear(channels, 8)
|
|
|
|
def __call__(self, x: SparseTensor, subdiv: Optional[SparseTensor] = None):
|
|
from .tensor import channel2spatial
|
|
|
|
if self.pred_subdiv:
|
|
subdiv = self.to_subdiv(x)
|
|
mask = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None
|
|
|
|
h = x.replace(self.norm1(x.feats))
|
|
h = h.replace(nn.silu(h.feats))
|
|
h = self.conv1(h)
|
|
h = channel2spatial(h, mask, 2)
|
|
xu = channel2spatial(x, mask, 2)
|
|
|
|
h = h.replace(self.norm2(h.feats))
|
|
h = h.replace(nn.silu(h.feats))
|
|
h = self.conv2(h)
|
|
|
|
reps = self.out_channels // max(self.channels // 8, 1)
|
|
skip = mx.repeat(xu.feats, reps, axis=1) if reps > 1 else xu.feats
|
|
out = h.replace(h.feats + skip)
|
|
return (out, subdiv) if self.pred_subdiv else out
|