formatting

test/modules/test_distribution_output.py

@@ -17,24 +17,24 @@ BATCH_SIZE = 32
 TOL = 0.3
 START_TOL_MULTIPLE = 1
 
+
 def inv_softplus(y: np.ndarray) -> np.ndarray:
     return np.log(np.exp(y) - 1)
 
-def maximum_likelihood_estimate_sgd(
-    distr_output: DistributionOutput,
-    samples: torch.Tensor,
-    init_biases: List[np.ndarray] = None,
-    num_epochs: int = 5,
-    learning_rate: float = 1e-2
-):
+
+def maximum_likelihood_estimate_sgd(distr_output: DistributionOutput,
+                                    samples: torch.Tensor,
+                                    init_biases: List[np.ndarray] = None,
+                                    num_epochs: int = 5,
+                                    learning_rate: float = 1e-2):
     distr_output.in_features = 1
     arg_proj = distr_output.get_args_proj()
-    
+
     if init_biases is not None:
         for param, bias in zip(arg_proj.proj, init_biases):
             nn.init.constant_(param.bias, bias)
-    
-    dummy_data = torch.ones((len(samples),1))
+
+    dummy_data = torch.ones((len(samples), 1))
     dataset = TensorDataset(dummy_data, samples)
     train_data = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
 
@@ -56,28 +56,27 @@ def maximum_likelihood_estimate_sgd(
             num_batches += 1
             cumulative_loss += loss.item()
         print("Epoch %s, loss: %s" % (e, cumulative_loss / num_batches))
-    
+
     if len(distr_args[0].shape) == 1:
         return [
-            param.detach().numpy() for param in arg_proj(torch.ones((1,1)))
+            param.detach().numpy() for param in arg_proj(torch.ones((1, 1)))
         ]
 
     return [
-        param[0].detach().numpy() for param in arg_proj(torch.ones((1,1)))
+        param[0].detach().numpy() for param in arg_proj(torch.ones((1, 1)))
     ]
 
 
-
 @pytest.mark.parametrize("df, loc, scale,", [(6.0, 2.3, 0.7)])
 def test_studentT_likelihood(df: float, loc: float, scale: float):
-    
-    dfs = torch.zeros((NUM_SAMPLES,)) + df
-    locs = torch.zeros((NUM_SAMPLES,)) + loc
-    scales = torch.zeros((NUM_SAMPLES,)) + scale
+
+    dfs = torch.zeros((NUM_SAMPLES, )) + df
+    locs = torch.zeros((NUM_SAMPLES, )) + loc
+    scales = torch.zeros((NUM_SAMPLES, )) + scale
 
     distr = StudentT(df=dfs, loc=locs, scale=scales)
     samples = distr.sample()
-    
+
     init_bias = [
         inv_softplus(df - 2),
         loc - START_TOL_MULTIPLE * TOL * loc,
@@ -89,15 +88,11 @@ def test_studentT_likelihood(df: float, loc: float, scale: float):
         samples,
         init_biases=init_bias,
         num_epochs=10,
-        learning_rate=1e-2
-    )
+        learning_rate=1e-2)
 
-    assert (
-        np.abs(df_hat - df) < TOL * df
-    ), f"df did not match: df = {df}, df_hat = {df_hat}"
-    assert (
-        np.abs(loc_hat - loc) < TOL * loc
-    ), f"loc did not match: loc = {loc}, loc_hat = {loc_hat}"
-    assert (
-        np.abs(scale_hat - scale) < TOL * scale
-    ), f"scale did not match: scale = {scale}, scale_hat = {scale_hat}"
+    assert (np.abs(df_hat - df) <
+            TOL * df), f"df did not match: df = {df}, df_hat = {df_hat}"
+    assert (np.abs(loc_hat - loc) <
+            TOL * loc), f"loc did not match: loc = {loc}, loc_hat = {loc_hat}"
+    assert (np.abs(scale_hat - scale) < TOL * scale
+            ), f"scale did not match: scale = {scale}, scale_hat = {scale_hat}"