fixed device testing so it works on a cpu

tests/integration_tests/scientific_machine_learning/test_feed_forward_methods.py

@@ -14,13 +14,12 @@ def test_device():
         sml.BayesianLinear(1, 1, device=cpu),
     )
     model = sml.FeedForwardNeuralNetwork(network)
-    device = (
-        torch.device("cuda", 0)
-        if torch.cuda.is_available()
-        else torch.device("mps", 0)
-        if torch.backends.mps.is_available()
-        else "cpu"
-    )
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
     model.to(device)
     assert model.network[0].weight.device == device
     assert model.network[1].weight_mu.device == device

tests/unit_tests/scientific_machine_learning/functional/test_functional_generalized_js_divergence.py

@@ -43,15 +43,13 @@ def test_device():
     """Note if neither cuda nor mps is available, this test will always pass"""
     posterior_distribution = [uq.Uniform(3, 1)]
     prior_distribution = [uq.Uniform(2, 5)]
-    device = (
-        torch.device("cuda", 0)
-        if torch.cuda.is_available()
-        else torch.device("mps", 0) if torch.backends.mps.is_available() else "cpu"
-    )
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
     divergence = func.generalized_jensen_shannon_divergence(
-        posterior_distribution,
-        prior_distribution,
-        alpha=0.0,
-        device=device
+        posterior_distribution, prior_distribution, alpha=0.0, device=device
     )
     assert divergence.device == device

tests/unit_tests/scientific_machine_learning/layers/test_bayesian_conv1d.py

@@ -54,11 +54,12 @@ def test_device():
     cpu = torch.device("cpu")
     layer = sml.BayesianConv1d(1, 1, 1, device=cpu)
     assert layer.weight_mu.device == cpu
-    device = (
-        torch.device("cuda", 0)
-        if torch.cuda.is_available()
-        else torch.device("mps", 0) if torch.backends.mps.is_available() else "cpu"
-    )
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
     layer.to(device)
     assert layer.weight_mu.device == device
 

tests/unit_tests/scientific_machine_learning/layers/test_bayesian_conv2d.py

@@ -106,11 +106,12 @@ def test_device():
     cpu = torch.device("cpu")
     layer = sml.BayesianConv2d(1, 1, 1, device=cpu)
     assert layer.weight_mu.device == cpu
-    device = (
-        torch.device("cuda", 0)
-        if torch.cuda.is_available()
-        else torch.device("mps", 0) if torch.backends.mps.is_available() else "cpu"
-    )
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
     layer.to(device)
     assert layer.weight_mu.device == device
 

tests/unit_tests/scientific_machine_learning/layers/test_bayesian_conv3d.py

@@ -38,7 +38,7 @@ def compute_hwl_out(
 )
 @settings(deadline=1_000)
 def test_default_output_shape(n, shape, in_channels, out_channels):
-    """Test the output """
+    """Test the output"""
     x = torch.rand((n, in_channels, *shape))
     layer = sml.BayesianConv3d(in_channels, out_channels, 1)
     y = layer(x)
@@ -97,6 +97,21 @@ def test_fancy_output_shape(kernel_size, stride, padding, dilation):
     assert y.shape == torch.Size([n, out_channels, h_out, w_out, l_out])
 
 
+def test_device():
+    """Note if neither cuda nor mps is available this test will always pass"""
+    cpu = torch.device("cpu")
+    layer = sml.BayesianConv3d(1, 1, 1, device=cpu)
+    assert layer.weight_mu.device == cpu
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
+    layer.to(device)
+    assert layer.weight_mu.device == device
+
+
 def test_deterministic_output():
     x = torch.rand(1, 1, 256, 256, 256)
     layer = sml.BayesianConv3d(1, 1, 1)

tests/unit_tests/scientific_machine_learning/layers/test_bayesian_fourier1d.py

@@ -23,11 +23,12 @@ def test_device():
     cpu = torch.device("cpu")
     layer = sml.BayesianFourier1d(1, 1, device=cpu)
     assert layer.weight_spectral_mu.device == cpu
-    device = (
-        torch.device("cuda", 0)
-        if torch.cuda.is_available()
-        else torch.device("mps", 0) if torch.backends.mps.is_available() else "cpu"
-    )
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
     layer.to(device)
     assert layer.weight_spectral_mu.device == device
 
@@ -66,4 +67,3 @@ def test_bias_false():
     layer = sml.BayesianFourier1d(1, 1, bias=False)
     y = layer(x)
     assert torch.all(y == torch.zeros_like(y))
-

tests/unit_tests/scientific_machine_learning/layers/test_bayesian_fourier_2d.py

@@ -23,6 +23,21 @@ def test_output_shape(batch_size, width, w, h, modes):
     assert x.shape == y.shape
 
 
+def test_device():
+    """Note if neither cuda nor mps is available this test will always pass"""
+    cpu = torch.device("cpu")
+    layer = sml.BayesianFourier2d(1, (1, 1), device=cpu)
+    assert layer.weight_spectral_mu.device == cpu
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
+    layer.to(device)
+    assert layer.weight_spectral_mu.device == device
+
+
 def test_deterministic_output():
     x = torch.rand((1, 1, 64, 128))
     layer = sml.BayesianFourier2d(1, (33, 65))

tests/unit_tests/scientific_machine_learning/layers/test_bayesian_fourier_3d.py

@@ -20,6 +20,21 @@ def test_output_shape(batch_size, width, d, w, h):
     assert x.shape == y.shape
 
 
+def test_device():
+    """Note if neither cuda nor mps is available this test will always pass"""
+    cpu = torch.device("cpu")
+    layer = sml.BayesianFourier3d(1, (1, 1, 1), device=cpu)
+    assert layer.weight_spectral_mu.device == cpu
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
+    layer.to(device)
+    assert layer.weight_spectral_mu.device == device
+
+
 def test_deterministic_output():
     x = torch.rand((1, 1, 32, 64, 128))
     layer = sml.BayesianFourier3d(1, (17, 33, 65))

tests/unit_tests/scientific_machine_learning/layers/test_bayesian_linear.py

@@ -22,11 +22,12 @@ def test_device():
     cpu = torch.device("cpu")
     layer = sml.BayesianLinear(1, 1, device=cpu)
     assert layer.weight_mu.device == cpu
-    device = (
-        torch.device("cuda", 0)
-        if torch.cuda.is_available()
-        else torch.device("mps", 0) if torch.backends.mps.is_available() else "cpu"
-    )
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
     layer.to(device)
     assert layer.weight_mu.device == device
 

tests/unit_tests/scientific_machine_learning/losses/test_gaussian_kl_divergence.py

@@ -29,11 +29,12 @@ def test_reduction_none():
 
 def test_device():
     """Note if neither cuda nor mps is available, this test will always pass"""
-    device = (
-        torch.device("cuda", 0)
-        if torch.cuda.is_available()
-        else torch.device("mps", 0) if torch.backends.mps.is_available() else "cpu"
-    )
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
     model = sml.FeedForwardNeuralNetwork(sml.BayesianLinear(1, 1))
     model.to(device)
     divergence_function = sml.GaussianKullbackLeiblerDivergence(device=device)

tests/unit_tests/scientific_machine_learning/losses/test_generalized_js_divergence.py

@@ -15,7 +15,9 @@ def test_reduction_shape():
         sml.BayesianLinear(width, 1),
     )
     model = sml.FeedForwardNeuralNetwork(network)
-    divergence_function = sml.GeneralizedJensenShannonDivergence(uq.Normal, uq.Uniform, n_samples=1)
+    divergence_function = sml.GeneralizedJensenShannonDivergence(
+        uq.Normal, uq.Uniform, n_samples=1
+    )
     divergence = divergence_function(model)
     assert divergence.shape == torch.Size()
 
@@ -26,11 +28,12 @@ def test_reduction_none_raises_error():
 
 
 def test_device():
-    device = (
-        torch.device("cuda", 0)
-        if torch.cuda.is_available()
-        else torch.device("mps", 0) if torch.backends.mps.is_available() else "cpu"
-    )
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
     divergence_function = sml.GeneralizedJensenShannonDivergence(
         uq.Normal, uq.Normal, device=device, n_samples=1
     )

tests/unit_tests/scientific_machine_learning/losses/test_geometric_js_divergence.py

@@ -28,11 +28,13 @@ def test_reduction_none():
 
 def test_device():
     """Note if neither cuda nor mps is available, this test will always pass"""
-    device = (
-        torch.device("cuda", 0)
-        if torch.cuda.is_available()
-        else torch.device("mps", 0) if torch.backends.mps.is_available() else "cpu"
-    )
+    if torch.cuda.is_available():
+        device = torch.device("cuda", 0)
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps", 0)
+    else:
+        device = torch.device("cpu")
+
     model = sml.FeedForwardNeuralNetwork(sml.BayesianLinear(1, 1))
     model.to(device)
     divergence_function = sml.GeometricJensenShannonDivergence(device=device)