diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..1489fea --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,27 @@ +import pytest +import torch + + +def available_devices(): + """Return every device the running machine can actually use. + + CPU is always present. Accelerators are appended only when their backend + reports as available, so the same test suite exercises CUDA, MPS, or XPU on + machines that have them while still running everywhere else. + """ + devices = ["cpu"] + if torch.cuda.is_available(): + devices.append("cuda") + if ( + getattr(torch.backends, "mps", None) is not None + and torch.backends.mps.is_available() + ): + devices.append("mps") + if getattr(torch, "xpu", None) is not None and torch.xpu.is_available(): + devices.append("xpu") + return devices + + +@pytest.fixture(params=available_devices()) +def device(request): + return torch.device(request.param) diff --git a/tests/test_cells.py b/tests/test_cells.py index 85f06ff..2c8ab9c 100755 --- a/tests/test_cells.py +++ b/tests/test_cells.py @@ -249,6 +249,41 @@ def test_cell_gradients(Cell, in_size, hid_size, _): assert p.grad is not None +@pytest.mark.parametrize("Cell, in_size, hid_size, double", CELL_CASES) +def test_cell_runs_on_device(Cell, in_size, hid_size, double, device): + """Every cell should forward, init state, and backward on each available + device (cpu plus any accelerator: cuda, mps, xpu).""" + cell = Cell(in_size, hid_size, bias=False).to(device) + + B = 4 + x = torch.randn(B, in_size, device=device, requires_grad=True) + + if double: + h, c = cell(x, (None, None)) + assert h.device.type == device.type + assert c.device.type == device.type + assert h.shape == (B, hid_size) + assert c.shape == (B, hid_size) + # explicit state already living on the device should be accepted + h2, c2 = cell(x, (h, c)) + assert h2.device.type == device.type + assert c2.device.type == device.type + out = h + else: + h = cell(x) + assert h.device.type == device.type + assert h.shape == (B, hid_size) + h2 = cell(x, h) + assert h2.device.type == device.type + out = h + + out.sum().backward() + for p in cell.parameters(): + if p.requires_grad: + assert p.grad is not None + assert p.grad.device.type == device.type + + @skip_windows @pytest.mark.parametrize("Cell, in_size, hid_size, double", CELL_CASES) def test_cell_compile(Cell, in_size, hid_size, double): diff --git a/tests/test_layers.py b/tests/test_layers.py index bee0445..7935257 100755 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -161,6 +161,45 @@ def test_layer_shapes_and_state(Layer, is_double): assert state_bf.shape == (num_layers, batch_size, hidden_size) +@pytest.mark.parametrize("Layer, is_double", LAYER_CASES) +def test_layer_runs_on_device(Layer, is_double, device): + """Every stacked layer should forward and backward on each available device + (cpu plus any accelerator: cuda, mps, xpu).""" + input_size, hidden_size = 5, 7 + seq_len, batch_size, num_layers = 4, 3, 2 + + layer = Layer( + input_size, + hidden_size, + num_layers=num_layers, + dropout=0.0, + batch_first=False, + bias=False, + ).to(device) + + x = torch.randn(seq_len, batch_size, input_size, device=device, requires_grad=True) + out, state = layer(x) + + assert out.device.type == device.type + assert out.shape == (seq_len, batch_size, hidden_size) + + if is_double: + h, c = state + assert h.device.type == device.type + assert c.device.type == device.type + assert h.shape == (num_layers, batch_size, hidden_size) + assert c.shape == (num_layers, batch_size, hidden_size) + else: + assert state.device.type == device.type + assert state.shape == (num_layers, batch_size, hidden_size) + + out.sum().backward() + for p in layer.parameters(): + if p.requires_grad: + assert p.grad is not None + assert p.grad.device.type == device.type + + @pytest.mark.parametrize("Layer", LAYER_CLASSES) def test_default_repr_shows_input_hidden(Layer): # Default repr should exactly match "Class(input_size, hidden_size)"