File size: 4,093 Bytes
ecc81b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """Device-placement conformance: everything that works on CPU works on an
accelerator, and CPU/device mismatches either work or fail loudly at the API
boundary — never three frames deep in an indexing kernel.
Runs against whatever accelerator the machine has — MPS on Apple Silicon,
CUDA elsewhere — and skips (visibly, per the suite's skips-are-not-passes
rule) when there is none. Device-placement bugs are indifferent to *which*
second device exists; what they need is any second device at all. What this
file cannot vouch for on MPS: CUDA-specific kernel numerics, torch.compile
backends, and float64 (unsupported on MPS — everything here is float32).
"""
import pytest
import torch
import torch_dimensions as td
from torch_dimensions.compose.kernel import axial_contract
if torch.cuda.is_available():
DEV = "cuda"
elif torch.backends.mps.is_available():
DEV = "mps"
else:
DEV = None
pytestmark = pytest.mark.skipif(DEV is None, reason="no accelerator (cuda or mps) available")
def _sparse_lattice():
valid = torch.tensor(
[[True, True, True, False], [True, False, True, True], [True, True, False, True]]
)
return td.Lattice(shape=(3, 4), names=("h", "w"), valid=valid, time=True)
def test_a_sparse_model_trains_on_the_device():
lat = _sparse_lattice()
torch.manual_seed(0)
model = td.LSTM(8, 3, lat).to(DEV)
x = torch.randn(2, 5, 3, 4, 8, device=DEV)
out = model(x)
assert out.device.type == DEV
out.pow(2).mean().backward()
assert all(p.grad is not None and p.grad.device.type == DEV for p in model.parameters())
def test_absent_cells_are_still_inert_on_the_device():
lat = _sparse_lattice()
torch.manual_seed(0)
model = td.LSTM(8, 3, lat).to(DEV)
x = torch.randn(2, 5, 3, 4, 8, device=DEV)
noise = torch.randn_like(x) * 1e3 * (~lat.mask().to(DEV)).to(x.dtype)
assert torch.equal(model(x), model(x + noise))
@pytest.mark.parametrize("lattice_on_device", [False, True])
def test_gather_scatter_round_trip_across_device_mismatches(lattice_on_device):
"""A CPU lattice must handle device tensors and a device lattice must
handle CPU tensors: `flat_idx` lives wherever `valid` lives, and indexing
must not require the caller to know that."""
lat = _sparse_lattice()
if lattice_on_device:
lat = lat.to(DEV)
for x_device in ("cpu", DEV):
x = torch.randn(2, 5, 3, 4, 8, device=x_device) * lat.mask().to(x_device)
g = lat.gather(x)
assert g.device.type == torch.device(x_device).type
assert torch.equal(lat.scatter(g), x)
def test_axial_contract_renormalizes_on_the_device():
lat = td.Lattice(shape=(2, 4), valid=torch.tensor([[1, 1, 0, 0], [1, 1, 1, 1]]).bool())
mask = lat.mask().to(torch.float32).to(DEV)
x = torch.ones(1, 2, 4, 3, device=DEV) * mask
row_stochastic = torch.full((4, 4), 0.25, device=DEV)
out = axial_contract(x, lat, 1, row_stochastic, valid=mask)
# constant-input invariant: a convex combination of ones is one at every
# present position, on any device
present = out.masked_select(mask.bool().expand_as(out))
assert torch.allclose(present, torch.ones_like(present), atol=1e-6)
def test_spec_survives_a_device_model():
lat = _sparse_lattice()
model = td.GRU(6, 3, lat).to(DEV)
spec = model.to_spec()
assert spec["lattice"]["cells"]["present"] == 9
assert spec["model"]["kind"] == "GRU"
@pytest.mark.parametrize("cls", [td.S4D, td.Mamba])
def test_the_ssm_family_runs_and_agrees_on_the_device(cls):
lat = _sparse_lattice()
torch.manual_seed(0)
cpu = cls(8, 3, lat)
dev = cls(8, 3, lat)
dev.load_state_dict(cpu.state_dict())
dev = dev.to(DEV)
x = torch.randn(2, 5, 3, 4, 8)
with torch.no_grad():
yc = cpu(x)
ym = dev(x.to(DEV))
assert not bool(ym.isnan().any())
diff = (yc - ym.cpu()).abs().max().item()
assert diff < 1e-4, f"cpu-vs-device diff {diff:.2e}"
dev(x.to(DEV)).pow(2).mean().backward()
assert all(p.grad is not None for p in dev.parameters() if p.requires_grad)
|