Refactor code for improved readability and consistency in tensor handling across multiple files

Author: Hananel-Hazan

bindsnet/evaluation/evaluation.py

@@ -44,7 +44,9 @@ def assign_labels(
             indices = torch.nonzero(labels == i).view(-1)
 
             # Compute average firing rates for this label.
-            selected_spikes = torch.index_select(spikes, dim=0, index=torch.tensor(indices))
+            selected_spikes = torch.index_select(
+                spikes, dim=0, index=torch.tensor(indices)
+            )
             rates[:, i] = alpha * rates[:, i] + (
                 torch.sum(selected_spikes, 0) / n_labeled
             )

bindsnet/learning/MCC_learning.py

@@ -103,7 +103,9 @@ def update(self, **kwargs) -> None:
             self, NoOp
         ):
             if self.feature_value.is_sparse:
-                self.feature_value = self.feature_value.to_dense().clamp_(self.min, self.max).to_sparse()
+                self.feature_value = (
+                    self.feature_value.to_dense().clamp_(self.min, self.max).to_sparse()
+                )
             else:
                 self.feature_value.clamp_(self.min, self.max)
 
@@ -252,8 +254,7 @@ def _connection_update(self, **kwargs) -> None:
             else:
                 if self.feature_value.is_sparse:
                     self.feature_value -= (
-                            torch.bmm(source_s, target_x)
-                            * self.connection.dt
+                        torch.bmm(source_s, target_x) * self.connection.dt
                     ).to_sparse()
                 else:
                     self.feature_value -= (
@@ -289,8 +290,7 @@ def _connection_update(self, **kwargs) -> None:
             else:
                 if self.feature_value.is_sparse:
                     self.feature_value += (
-                            torch.bmm(source_x, target_s)
-                            * self.connection.dt
+                        torch.bmm(source_x, target_s) * self.connection.dt
                     ).to_sparse()
                 else:
                     self.feature_value += (
@@ -524,9 +524,7 @@ def _connection_update(self, **kwargs) -> None:
             ) % self.average_update
 
             if self.continues_update or self.average_buffer_index == 0:
-                update = self.nu[0] * torch.mean(
-                    self.average_buffer, dim=0
-                )
+                update = self.nu[0] * torch.mean(self.average_buffer, dim=0)
                 if self.feature_value.is_sparse:
                     update = update.to_sparse()
                 self.feature_value += update

bindsnet/models/models.py

@@ -10,7 +10,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
 
 
@@ -182,17 +186,17 @@ def __init__(
             device=device,
             pipeline=[
                 Weight(
-                    'weight',
+                    "weight",
                     w,
                     range=[wmin, wmax],
                     norm=norm,
                     reduction=reduction,
                     nu=nu,
                     learning_rule=MMCPostPre,
                     sparse=sparse,
-                    batch_size=batch_size
+                    batch_size=batch_size,
                 )
-            ]
+            ],
         )
         w = self.exc * torch.diag(torch.ones(self.n_neurons))
         if sparse:
@@ -201,14 +205,7 @@ def __init__(
             source=exc_layer,
             target=inh_layer,
             device=device,
-            pipeline=[
-                Weight(
-                    'weight',
-                    w,
-                    range=[0, self.exc],
-                    sparse=sparse
-                )
-            ]
+            pipeline=[Weight("weight", w, range=[0, self.exc], sparse=sparse)],
         )
         w = -self.inh * (
             torch.ones(self.n_neurons, self.n_neurons)
@@ -220,14 +217,7 @@ def __init__(
             source=inh_layer,
             target=exc_layer,
             device=device,
-            pipeline=[
-                Weight(
-                    'weight',
-                    w,
-                    range=[-self.inh, 0],
-                    sparse=sparse
-                )
-            ]
+            pipeline=[Weight("weight", w, range=[-self.inh, 0], sparse=sparse)],
         )
 
         # Add to network

bindsnet/network/monitors.py

@@ -45,7 +45,7 @@ def __init__(
         time: Optional[int] = None,
         batch_size: int = 1,
         device: str = "cpu",
-        sparse: Optional[bool] = False
+        sparse: Optional[bool] = False,
     ):
         # language=rst
         """
@@ -100,9 +100,9 @@ def record(self) -> None:
         for v in self.state_vars:
             data = getattr(self.obj, v).unsqueeze(0)
             # self.recording[v].append(data.detach().clone().to(self.device))
-            record = torch.empty_like(data, device=self.device, requires_grad=False).copy_(
-                data, non_blocking=True
-            )
+            record = torch.empty_like(
+                data, device=self.device, requires_grad=False
+            ).copy_(data, non_blocking=True)
             if self.sparse:
                 record = record.to_sparse()
             self.recording[v].append(record)

bindsnet/network/topology_features.py

@@ -134,8 +134,9 @@ def __init__(
             self.value = self.value.unsqueeze(0).repeat(self.batch_size, 1, 1)
 
         self.value = self.value.to_sparse()
-        assert not getattr(self, 'enforce_polarity', False), \
-            "enforce_polarity isn't supported for sparse tensors"
+        assert not getattr(
+            self, "enforce_polarity", False
+        ), "enforce_polarity isn't supported for sparse tensors"
 
     @abstractmethod
     def reset_state_variables(self) -> None:
@@ -179,9 +180,15 @@ def prime_feature(self, connection, device, **kwargs) -> None:
         # Check if values/norms are the correct shape
         if isinstance(self.value, torch.Tensor):
             if self.sparse:
-                assert tuple(self.value.shape[1:]) == (connection.source.n, connection.target.n)
+                assert tuple(self.value.shape[1:]) == (
+                    connection.source.n,
+                    connection.target.n,
+                )
             else:
-                assert tuple(self.value.shape) == (connection.source.n, connection.target.n)
+                assert tuple(self.value.shape) == (
+                    connection.source.n,
+                    connection.target.n,
+                )
 
         if self.norm is not None and isinstance(self.norm, torch.Tensor):
             assert self.norm.shape[0] == connection.target.n
@@ -325,10 +332,12 @@ def assert_feature_in_range(self):
 
     def assert_valid_shape(self, source_shape, target_shape, f):
         # Multidimensional feat
-        if (not self.sparse and len(f.shape) > 1) or (self.sparse and len(f.shape[1:]) > 1):
+        if (not self.sparse and len(f.shape) > 1) or (
+            self.sparse and len(f.shape[1:]) > 1
+        ):
             if self.sparse:
                 f_shape = f.shape[1:]
-                expected = ('batch_size', source_shape, target_shape)
+                expected = ("batch_size", source_shape, target_shape)
             else:
                 f_shape = f.shape
                 expected = (source_shape, target_shape)
@@ -352,7 +361,7 @@ def __init__(
         decay: float = 0.0,
         parent_feature=None,
         sparse: Optional[bool] = False,
-        batch_size: int = 1
+        batch_size: int = 1,
     ) -> None:
         # language=rst
         """
@@ -386,19 +395,15 @@ def __init__(
             decay=decay,
             parent_feature=parent_feature,
             sparse=sparse,
-            batch_size=batch_size
+            batch_size=batch_size,
         )
 
     def sparse_bernoulli(self):
         values = torch.bernoulli(self.value.values())
         mask = values != 0
         indices = self.value.indices()[:, mask]
         non_zero = values[mask]
-        return torch.sparse_coo_tensor(
-            indices,
-            non_zero,
-            self.value.size()
-        )
+        return torch.sparse_coo_tensor(indices, non_zero, self.value.size())
 
     def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
         if self.sparse:
@@ -448,7 +453,7 @@ def __init__(
         name: str,
         value: Union[torch.Tensor, float, int] = None,
         sparse: Optional[bool] = False,
-        batch_size: int = 1
+        batch_size: int = 1,
     ) -> None:
         # language=rst
         """
@@ -472,12 +477,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,
-            sparse=sparse,
-            batch_size=batch_size
-        )
+        super().__init__(name=name, value=value, sparse=sparse, batch_size=batch_size)
 
     def compute(self, conn_spikes) -> torch.Tensor:
         return conn_spikes * self.value
@@ -561,7 +561,7 @@ def __init__(
         enforce_polarity: Optional[bool] = False,
         decay: float = 0.0,
         sparse: Optional[bool] = False,
-        batch_size: int = 1
+        batch_size: int = 1,
     ) -> None:
         # language=rst
         """
@@ -596,7 +596,7 @@ def __init__(
             reduction=reduction,
             decay=decay,
             sparse=sparse,
-            batch_size=batch_size
+            batch_size=batch_size,
         )
 
     def reset_state_variables(self) -> None:
@@ -651,7 +651,7 @@ def __init__(
         range: Optional[Sequence[float]] = None,
         norm: Optional[Union[torch.Tensor, float, int]] = None,
         sparse: Optional[bool] = False,
-        batch_size: int = 1
+        batch_size: int = 1,
     ) -> None:
         # language=rst
         """
@@ -671,7 +671,7 @@ def __init__(
             range=[-torch.inf, +torch.inf] if range is None else range,
             norm=norm,
             sparse=sparse,
-            batch_size=batch_size
+            batch_size=batch_size,
         )
 
     def reset_state_variables(self) -> None:
@@ -697,7 +697,7 @@ def __init__(
         value: Union[torch.Tensor, float, int] = None,
         range: Optional[Sequence[float]] = None,
         sparse: Optional[bool] = False,
-        batch_size: int = 1
+        batch_size: int = 1,
     ) -> None:
         # language=rst
         """
@@ -708,7 +708,9 @@ def __init__(
         :param batch_size: Mini-batch size.
         """
 
-        super().__init__(name=name, value=value, range=range, sparse=sparse, batch_size=batch_size)
+        super().__init__(
+            name=name, value=value, range=range, sparse=sparse, batch_size=batch_size
+        )
 
     def reset_state_variables(self) -> None:
         pass
@@ -738,7 +740,7 @@ def __init__(
         degrade_function: callable = None,
         parent_feature: Optional[AbstractFeature] = None,
         sparse: Optional[bool] = False,
-        batch_size: int = 1
+        batch_size: int = 1,
     ) -> None:
         # language=rst
         """
@@ -754,7 +756,13 @@ def __init__(
         """
 
         # Note: parent_feature will override value. See abstract constructor
-        super().__init__(name=name, value=value, parent_feature=parent_feature, sparse=sparse, batch_size=batch_size)
+        super().__init__(
+            name=name,
+            value=value,
+            parent_feature=parent_feature,
+            sparse=sparse,
+            batch_size=batch_size,
+        )
 
         self.degrade_function = degrade_function
 
@@ -774,7 +782,7 @@ def __init__(
         const_update_rate: float = 0.1,
         const_decay: float = 0.001,
         sparse: Optional[bool] = False,
-        batch_size: int = 1
+        batch_size: int = 1,
     ) -> None:
         # language=rst
         """
@@ -833,7 +841,9 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
                 flat_conn_spikes = conn_spikes.to_dense().flatten()
             else:
                 flat_conn_spikes = conn_spikes.flatten()
-            self.spike_buffer[:, self.counter % self.spike_buffer.shape[1]] = flat_conn_spikes
+            self.spike_buffer[:, self.counter % self.spike_buffer.shape[1]] = (
+                flat_conn_spikes
+            )
             self.counter += 1
 
         # Update the masks
@@ -872,7 +882,7 @@ def __init__(
         const_update_rate: float = 0.1,
         const_decay: float = 0.01,
         sparse: Optional[bool] = False,
-        batch_size: int = 1
+        batch_size: int = 1,
     ) -> None:
         # language=rst
         """
@@ -931,7 +941,9 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
                 flat_conn_spikes = conn_spikes.to_dense().flatten()
             else:
                 flat_conn_spikes = conn_spikes.flatten()
-            self.spike_buffer[:, self.counter % self.spike_buffer.shape[1]] = flat_conn_spikes
+            self.spike_buffer[:, self.counter % self.spike_buffer.shape[1]] = (
+                flat_conn_spikes
+            )
             self.counter += 1
 
         # Update the masks