95fd4496da
All verified against torch, which IS available here - only spconv was not. So unlike the submanifold conv these compare against upstream's real semantics rather than a hand-written paraphrase. 6/6 pass, max err 9.5e-7. The norm distinction is the trap: LayerNorm32 is applied to x.feats (per-voxel over channels) while SparseGroupNorm32 reshapes [N_b,C] -> [1,C,N_b] per batch item, so its statistics span channels-in-group AND voxels. Both produce identical shapes, so a mixup is numerically silent - there is an explicit test asserting the group norm does NOT match a per-voxel group_norm. Architecture confirmed from the converted weights rather than constructor defaults: no rope, no qk_rms_norm, transformer norms non-affine, ResBlock norm1 affine/norm2 not.
253 lines
9.8 KiB
Python
253 lines
9.8 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):
|
|
def __init__(self, channels: int, mlp_ratio: float = 4.0):
|
|
super().__init__()
|
|
hidden = int(channels * mlp_ratio)
|
|
self.mlp = [nn.Linear(channels, hidden), None, nn.Linear(hidden, channels)]
|
|
|
|
def __call__(self, x: SparseTensor) -> SparseTensor:
|
|
h = self.mlp[0](x.feats)
|
|
h = nn.gelu_approx(h)
|
|
return x.replace(self.mlp[2](h))
|
|
|
|
|
|
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)
|