Refactor code for improved readability and consistency in weight handling and warnings

Author: Hananel-Hazan

bindsnet/models/models.py

@@ -9,7 +9,11 @@
 from bindsnet.learning.MCC_learning import PostPre as MMCPostPre
 from bindsnet.network import Network
 from bindsnet.network.nodes import DiehlAndCookNodes, Input, LIFNodes
-from bindsnet.network.topology import Connection, LocalConnection, MulticompartmentConnection
+from bindsnet.network.topology import (
+    Connection,
+    LocalConnection,
+    MulticompartmentConnection,
+)
 from bindsnet.network.topology_features import Weight
 
 
@@ -181,30 +185,23 @@ def __init__(
             device=device,
             pipeline=[
                 Weight(
-                    'weight',
+                    "weight",
                     w,
                     value_dtype=w_dtype,
                     range=[wmin, wmax],
                     norm=norm,
                     reduction=reduction,
                     nu=nu,
-                    learning_rule=MMCPostPre
+                    learning_rule=MMCPostPre,
                 )
-            ]
+            ],
         )
         w = self.exc * torch.diag(torch.ones(self.n_neurons))
         exc_inh_conn = MulticompartmentConnection(
             source=exc_layer,
             target=inh_layer,
             device=device,
-            pipeline=[
-                Weight(
-                    'weight',
-                    w,
-                    value_dtype=w_dtype,
-                    range=[0, self.exc]
-                )
-            ]
+            pipeline=[Weight("weight", w, value_dtype=w_dtype, range=[0, self.exc])],
         )
         w = -self.inh * (
             torch.ones(self.n_neurons, self.n_neurons)
@@ -214,14 +211,7 @@ def __init__(
             source=inh_layer,
             target=exc_layer,
             device=device,
-            pipeline=[
-                Weight(
-                    'weight',
-                    w,
-                    value_dtype=w_dtype,
-                    range=[-self.inh, 0]
-                )
-            ]
+            pipeline=[Weight("weight", w, value_dtype=w_dtype, range=[-self.inh, 0])],
         )
 
         # Add to network

bindsnet/network/topology.py

@@ -144,7 +144,9 @@ def reset_state_variables(self) -> None:
     @staticmethod
     def cast_dtype_if_needed(w, w_dtype):
         if w.dtype != w_dtype:
-            warnings.warn(f"Provided w has data type {w.dtype} but parameter w_dtype is {w_dtype}")
+            warnings.warn(
+                f"Provided w has data type {w.dtype} but parameter w_dtype is {w_dtype}"
+            )
             return w.to(dtype=w_dtype)
         else:
             return w

bindsnet/network/topology_features.py

@@ -127,7 +127,9 @@ def __init__(
     @staticmethod
     def cast_dtype_if_needed(value, value_dtype):
         if value.dtype != value_dtype:
-            warnings.warn(f"Provided value has data type {value.dtype} but parameter w_dtype is {value_dtype}")
+            warnings.warn(
+                f"Provided value has data type {value.dtype} but parameter w_dtype is {value_dtype}"
+            )
             return value.to(dtype=value_dtype)
         else:
             return value
@@ -431,11 +433,7 @@ def __init__(
             # Send boolean to tensor (priming wont work if it's not a tensor)
             value = torch.tensor(value)
 
-        super().__init__(
-            name=name,
-            value=value,
-            value_dtype=torch.bool
-        )
+        super().__init__(name=name, value=value, value_dtype=torch.bool)
 
         self.name = name
         self.value = value
@@ -705,7 +703,12 @@ def __init__(
         """
 
         # Note: parent_feature will override value. See abstract constructor
-        super().__init__(name=name, value=value, value_dtype=value_dtype, parent_feature=parent_feature)
+        super().__init__(
+            name=name,
+            value=value,
+            value_dtype=value_dtype,
+            parent_feature=parent_feature,
+        )
 
         self.degrade_function = degrade_function
 

examples/mnist/batch_eth_mnist.py

@@ -43,8 +43,8 @@
 parser.add_argument(
     "--w_dtype",
     type=str,
-    default='float32',
-    help='Datatype to use for weights. Examples: float32, float16, bfloat16 etc'
+    default="float32",
+    help="Datatype to use for weights. Examples: float32, float16, bfloat16 etc",
 )
 parser.add_argument("--train", dest="train", action="store_true")
 parser.add_argument("--test", dest="train", action="store_false")
@@ -109,7 +109,7 @@
     theta_plus=theta_plus,
     inpt_shape=(1, 28, 28),
     device=device,
-    w_dtype=getattr(torch, args.w_dtype)
+    w_dtype=getattr(torch, args.w_dtype),
 )
 
 # Directs network to GPU
@@ -279,7 +279,9 @@
             image = batch["image"][:, 0].view(28, 28)
             inpt = inputs["X"][:, 0].view(time, 784).sum(0).view(28, 28)
             lable = batch["label"][0]
-            input_exc_weights = network.connections[("X", "Ae")].feature_index['weight'].value
+            input_exc_weights = (
+                network.connections[("X", "Ae")].feature_index["weight"].value
+            )
             square_weights = get_square_weights(
                 input_exc_weights.view(784, n_neurons), n_sqrt, 28
             )