Sparse layer set: linear, norms, activations, resblock, attention, transformer blocks
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.
This commit is contained in:
John
parent
db2097c19f
commit
95fd4496da
252
lato_mlx/sparse/ops.py
Normal file
252
lato_mlx/sparse/ops.py
Normal file
@ -0,0 +1,252 @@
|
|||||||
|
"""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)
|
||||||
171
tests/test_ops.py
Normal file
171
tests/test_ops.py
Normal file
@ -0,0 +1,171 @@
|
|||||||
|
"""Sparse layer tests against torch.
|
||||||
|
|
||||||
|
Unlike the submanifold conv — where spconv is uninstallable and the oracle had to be
|
||||||
|
hand-written — every layer here has a real torch counterpart, so these compare against
|
||||||
|
upstream's actual semantics rather than a paraphrase of them. The upstream forward
|
||||||
|
bodies are reproduced verbatim (see modules/sparse/{norm,linear,nonlinearity}.py),
|
||||||
|
including the [N_b,C] -> [1,C,N_b] GroupNorm reshape, which is the one that would fail
|
||||||
|
silently if guessed.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import mlx.core as mx
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||||
|
|
||||||
|
from lato_mlx.sparse.ops import ( # noqa: E402
|
||||||
|
LayerNorm32,
|
||||||
|
SparseGroupNorm32,
|
||||||
|
SparseMultiHeadAttention,
|
||||||
|
SparseTransformerBlock,
|
||||||
|
)
|
||||||
|
from lato_mlx.sparse.tensor import SparseTensor # noqa: E402
|
||||||
|
|
||||||
|
|
||||||
|
def make_batched(n_per_batch, channels, seed=0):
|
||||||
|
"""Batch-contiguous coords, as upstream requires."""
|
||||||
|
rng = np.random.default_rng(seed)
|
||||||
|
coords, feats = [], []
|
||||||
|
for b, nb in enumerate(n_per_batch):
|
||||||
|
for i in range(nb):
|
||||||
|
coords.append((b, i // 16, (i // 4) % 4, i % 4))
|
||||||
|
feats.append(rng.standard_normal((nb, channels)).astype(np.float32))
|
||||||
|
return np.array(coords, dtype=np.int32), np.concatenate(feats, 0)
|
||||||
|
|
||||||
|
|
||||||
|
def test_group_norm_matches_torch(groups=8, channels=32):
|
||||||
|
n_per_batch = [37, 51]
|
||||||
|
coords, feats = make_batched(n_per_batch, channels, seed=1)
|
||||||
|
rng = np.random.default_rng(2)
|
||||||
|
w = rng.standard_normal(channels).astype(np.float32)
|
||||||
|
b = rng.standard_normal(channels).astype(np.float32)
|
||||||
|
|
||||||
|
gn = SparseGroupNorm32(groups, channels)
|
||||||
|
gn.weight, gn.bias = mx.array(w), mx.array(b)
|
||||||
|
got = np.asarray(gn(SparseTensor(mx.array(feats), mx.array(coords))).feats)
|
||||||
|
|
||||||
|
# upstream: per batch item, [N_b,C] -> permute -> [1,C,N_b] -> nn.GroupNorm
|
||||||
|
tg = torch.nn.GroupNorm(groups, channels, eps=1e-5, affine=True)
|
||||||
|
tg.weight.data = torch.tensor(w)
|
||||||
|
tg.bias.data = torch.tensor(b)
|
||||||
|
want = np.zeros_like(feats)
|
||||||
|
off = 0
|
||||||
|
for nb in n_per_batch:
|
||||||
|
bf = torch.tensor(feats[off : off + nb])
|
||||||
|
bf = bf.permute(1, 0).reshape(1, channels, -1)
|
||||||
|
bf = tg(bf)
|
||||||
|
want[off : off + nb] = bf.reshape(channels, -1).permute(1, 0).detach().numpy()
|
||||||
|
off += nb
|
||||||
|
|
||||||
|
err = np.abs(got - want).max()
|
||||||
|
assert err < 2e-4, f"group norm err {err:.3g}"
|
||||||
|
return err
|
||||||
|
|
||||||
|
|
||||||
|
def test_group_norm_is_not_per_voxel(groups=8, channels=32):
|
||||||
|
"""Guard the easy-to-miss distinction: GroupNorm here is NOT a per-voxel norm."""
|
||||||
|
coords, feats = make_batched([40], channels, seed=5)
|
||||||
|
gn = SparseGroupNorm32(groups, channels)
|
||||||
|
got = np.asarray(gn(SparseTensor(mx.array(feats), mx.array(coords))).feats)
|
||||||
|
per_voxel = torch.nn.functional.group_norm(
|
||||||
|
torch.tensor(feats).reshape(40, channels), groups
|
||||||
|
).numpy()
|
||||||
|
assert np.abs(got - per_voxel).max() > 1e-3, "matched per-voxel norm — reshape lost"
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
|
||||||
|
def test_layer_norm_matches_torch(channels=64):
|
||||||
|
rng = np.random.default_rng(3)
|
||||||
|
feats = rng.standard_normal((50, channels)).astype(np.float32)
|
||||||
|
ln = LayerNorm32(channels, affine=False, eps=1e-6)
|
||||||
|
got = np.asarray(ln(mx.array(feats)))
|
||||||
|
want = F.layer_norm(torch.tensor(feats), (channels,), eps=1e-6).numpy()
|
||||||
|
err = np.abs(got - want).max()
|
||||||
|
assert err < 1e-5, f"layer norm err {err:.3g}"
|
||||||
|
return err
|
||||||
|
|
||||||
|
|
||||||
|
def test_self_attention_matches_torch(channels=64, heads=8):
|
||||||
|
n_per_batch = [23, 31]
|
||||||
|
coords, feats = make_batched(n_per_batch, channels, seed=4)
|
||||||
|
rng = np.random.default_rng(6)
|
||||||
|
wq = rng.standard_normal((channels * 3, channels)).astype(np.float32) * 0.05
|
||||||
|
bq = rng.standard_normal((channels * 3,)).astype(np.float32) * 0.05
|
||||||
|
wo = rng.standard_normal((channels, channels)).astype(np.float32) * 0.05
|
||||||
|
bo = rng.standard_normal((channels,)).astype(np.float32) * 0.05
|
||||||
|
|
||||||
|
attn = SparseMultiHeadAttention(channels, heads)
|
||||||
|
attn.to_qkv.weight, attn.to_qkv.bias = mx.array(wq), mx.array(bq)
|
||||||
|
attn.to_out.weight, attn.to_out.bias = mx.array(wo), mx.array(bo)
|
||||||
|
got = np.asarray(attn(SparseTensor(mx.array(feats), mx.array(coords))).feats)
|
||||||
|
|
||||||
|
# reference: attention strictly within each batch item
|
||||||
|
d = channels // heads
|
||||||
|
want = np.zeros_like(feats)
|
||||||
|
off = 0
|
||||||
|
for nb in n_per_batch:
|
||||||
|
f = torch.tensor(feats[off : off + nb])
|
||||||
|
qkv = F.linear(f, torch.tensor(wq), torch.tensor(bq)).reshape(nb, 3, heads, d)
|
||||||
|
q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]
|
||||||
|
o = F.scaled_dot_product_attention(
|
||||||
|
q.permute(1, 0, 2)[None], k.permute(1, 0, 2)[None], v.permute(1, 0, 2)[None]
|
||||||
|
)
|
||||||
|
o = o[0].permute(1, 0, 2).reshape(nb, channels)
|
||||||
|
want[off : off + nb] = (
|
||||||
|
F.linear(o, torch.tensor(wo), torch.tensor(bo)).detach().numpy()
|
||||||
|
)
|
||||||
|
off += nb
|
||||||
|
|
||||||
|
err = np.abs(got - want).max()
|
||||||
|
assert err < 2e-4, f"attention err {err:.3g}"
|
||||||
|
return err
|
||||||
|
|
||||||
|
|
||||||
|
def test_attention_does_not_cross_batches(channels=32, heads=4):
|
||||||
|
"""Perturbing batch 1 must never change batch 0's output."""
|
||||||
|
coords, feats = make_batched([12, 12], channels, seed=7)
|
||||||
|
attn = SparseMultiHeadAttention(channels, heads)
|
||||||
|
a = np.asarray(attn(SparseTensor(mx.array(feats), mx.array(coords))).feats)
|
||||||
|
f2 = feats.copy()
|
||||||
|
f2[12:] += 10.0
|
||||||
|
b = np.asarray(attn(SparseTensor(mx.array(f2), mx.array(coords))).feats)
|
||||||
|
assert np.abs(a[:12] - b[:12]).max() < 1e-5, "batch 0 changed — attention leaked"
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
|
||||||
|
def test_transformer_block_shape(channels=64, heads=8):
|
||||||
|
coords, feats = make_batched([20, 20], channels, seed=8)
|
||||||
|
blk = SparseTransformerBlock(channels, heads)
|
||||||
|
out = blk(SparseTensor(mx.array(feats), mx.array(coords)))
|
||||||
|
assert out.feats.shape == (40, channels)
|
||||||
|
assert np.isfinite(np.asarray(out.feats)).all(), "non-finite output"
|
||||||
|
return 0.0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
tests = [
|
||||||
|
("group norm vs torch", test_group_norm_matches_torch),
|
||||||
|
("group norm != per-voxel", test_group_norm_is_not_per_voxel),
|
||||||
|
("layer norm vs torch", test_layer_norm_matches_torch),
|
||||||
|
("self-attn vs torch", test_self_attention_matches_torch),
|
||||||
|
("attn batch isolation", test_attention_does_not_cross_batches),
|
||||||
|
("transformer block", test_transformer_block_shape),
|
||||||
|
]
|
||||||
|
failed = 0
|
||||||
|
for name, fn in tests:
|
||||||
|
try:
|
||||||
|
err = fn()
|
||||||
|
print(f" PASS {name:26s} (max err {err:.2e})")
|
||||||
|
except AssertionError as e:
|
||||||
|
print(f" FAIL {name:26s} {e}")
|
||||||
|
failed += 1
|
||||||
|
except Exception as e: # noqa: BLE001
|
||||||
|
print(f" ERROR {name:26s} {type(e).__name__}: {e}")
|
||||||
|
failed += 1
|
||||||
|
print(f"\n{len(tests)-failed}/{len(tests)} passed")
|
||||||
|
sys.exit(1 if failed else 0)
|
||||||
Loading…
Reference in New Issue
Block a user