diff --git a/CHANGELOG.md b/CHANGELOG.md index dd9d6185..3f49c077 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,11 @@ see the [GitHub releases / tags](https://github.com/BindsNET/bindsnet/releases). - The CUDA `test_transfer` tests in `test/network/test_nodes.py` and `test/network/test_connections.py` are reported as skipped on machines without a GPU; before, they returned early and were reported as passed. +- A connection built from weights already on a GPU (`w=` on CUDA or XPU) no longer + fails in its constructor with "Expected all tensors to be on the same device". + `wmin` and `wmax` are now created on the device of the given `w`; before, they were + always on the CPU. Single-value and per-weight bounds both work. Results do not + change. ## [0.3.5] - 2026-10-29 diff --git a/bindsnet/network/topology.py b/bindsnet/network/topology.py index 3a009fa0..409c4953 100644 --- a/bindsnet/network/topology.py +++ b/bindsnet/network/topology.py @@ -87,13 +87,23 @@ def __init__( self.update_rule = kwargs.get("update_rule", None) + # The bounds live on the device of given weights: subclasses clamp w against + # them in their constructors, and a GPU w with CPU bounds raised a device + # mismatch. Without w the weights start on the CPU, and so do the bounds. + w_given = kwargs.get("w", None) + bound_device = w_given.device if isinstance(w_given, torch.Tensor) else None + # Float32 necessary for comparisons with +/-inf self.wmin = Parameter( - torch.as_tensor(kwargs.get("wmin", -np.inf), dtype=torch.float32), + torch.as_tensor( + kwargs.get("wmin", -np.inf), dtype=torch.float32, device=bound_device + ), requires_grad=False, ) self.wmax = Parameter( - torch.as_tensor(kwargs.get("wmax", np.inf), dtype=torch.float32), + torch.as_tensor( + kwargs.get("wmax", np.inf), dtype=torch.float32, device=bound_device + ), requires_grad=False, ) self.norm = kwargs.get("norm", None) diff --git a/test/network/test_connections.py b/test/network/test_connections.py index 3db015a1..fdb21c78 100644 --- a/test/network/test_connections.py +++ b/test/network/test_connections.py @@ -932,3 +932,48 @@ def test_windows_and_maxpool_state_stay_on_gpu(): net.run(inputs=inputs, time=5) net.reset_state_variables() assert pool.firing_rates.is_cuda + + +def test_weight_bounds_without_given_weights_stay_on_cpu(): + """Without w the weights start on the CPU, and so do the bounds.""" + conn = Connection(Input(n=4), LIFNodes(n=3), wmin=0.0, wmax=1.0) + assert conn.wmin.device.type == "cpu" + assert conn.wmax.device.type == "cpu" + assert conn.w.device.type == "cpu" + assert ((conn.w >= 0) & (conn.w <= 1)).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.parametrize( + "wmin, wmax", + [(0.0, 1.0), (torch.zeros(4, 3), torch.ones(4, 3))], + ids=["scalar-bounds", "per-weight-bounds"], +) +def test_connection_from_gpu_weights_clamps_and_learns_on_gpu(wmin, wmax): + """A Connection built from CUDA weights clamps them at construction and keeps + clamping during learning, all on the GPU. + + Caught: the constructor's torch.clamp(w, wmin, wmax) raised a device mismatch + because the bounds were on the CPU. + """ + w = torch.rand(4, 3, device="cuda") * 3 - 1 # values outside [0, 1] + conn = Connection( + Input(n=4, traces=True), + LIFNodes(n=3, traces=True), + w=w, + wmin=wmin, + wmax=wmax, + update_rule=PostPre, + nu=(1.0, 1.0), + ) + assert conn.w.is_cuda and conn.wmin.is_cuda and conn.wmax.is_cuda + assert ((conn.w >= 0) & (conn.w <= 1)).all() + + net = Network(dt=1.0) + net.add_layer(conn.source, "in") + net.add_layer(conn.target, "out") + net.add_connection(conn, "in", "out") + net.to("cuda") + net.run(inputs={"in": (torch.rand(20, 1, 4) < 0.5).cuda()}, time=20) + assert conn.w.is_cuda + assert ((conn.w >= 0) & (conn.w <= 1)).all()