tests: add gpu tests

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)

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):

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)"