sinabs 3.0.4.dev25__py3-none-any.whl → 3.1.1__py3-none-any.whl

sinabs/backend/dynapcnn/weight_rescaling_methods.py

@@ -0,0 +1,53 @@
+import statistics
+from typing import Iterable
+
+import numpy as np
+
+
+def rescale_method_1(scaling_factors: Iterable[int], lambda_: float = 0.5) -> float:
+    """
+        This method will use the average (scaled by `lambda_`) of the computed re-scaling factor
+    for the pooling layer(s) feeding into a convolutional layer.
+
+    Arguments
+    ---------
+    - scaling_factors (list): the list of re-scaling factors computed by each `SumPool2d` layer targeting a
+        single `Conv2d` layer within a `DynapcnnLayer` instance.
+    - lambda_ (float): a scaling variable that multiplies the computed average re-scaling factor of the pooling layers.
+
+    Returns
+    ---------
+    - the averaged re-scaling factor multiplied by `lambda_` if `len(scaling_factors) > 0`, else `1` is returned.
+    """
+
+    if len(scaling_factors) > 0:
+        return np.round(np.mean(list(scaling_factors)) * lambda_, 2)
+    else:
+        return 1.0
+
+
+def rescale_method_2(scaling_factors: Iterable[int], lambda_: float = 0.5) -> float:
+    """
+        This method will use the harmonic mean (scaled by `lambda_`) of the computed re-scaling factor
+    for the pooling layer(s) feeding into a convolutional layer.
+
+    Arguments
+    ---------
+    - scaling_factors (list): the list of re-scaling factors computed by each `SumPool2d` layer targeting a
+        single `Conv2d` layer within a `DynapcnnLayer` instance.
+    - lambda_ (float): a scaling variable that multiplies the computed average re-scaling factor of the pooling layers.
+
+    Returns
+    ---------
+    - the averaged re-scaling factor multiplied by `lambda_` if `len(scaling_factors) > 0`, else `1` is returned.
+
+    Note
+    ---------
+    - since the harmonic mean is less sensitive to outliers it **could be** that this is a better method
+        for weight re-scaling when multiple poolings with big differentces in kernel sizes are being considered.
+    """
+
+    if len(scaling_factors) > 0:
+        return np.round(statistics.harmonic_mean(list(scaling_factors)) * lambda_, 2)
+    else:
+        return 1.0

sinabs/conversion.py

@@ -9,7 +9,7 @@ def replace_module(model: nn.Module, source_class: type, mapper_fn: Callable):
     """A utility function that returns a copy of the model, where specific layers are replaced with
     another type depending on the mapper function.
 
-    Parameters:
+    Args:
         model: A PyTorch model.
         source_class: the layer class to replace. Each find will be passed to mapper_fn
         mapper_fn: A callable that takes as argument the layer to replace and returns the new object.
@@ -31,7 +31,7 @@ def replace_module_(model: nn.Sequential, source_class: type, mapper_fn: Callabl
     """In-place version of replace_module that will step through modules that have children and
     apply the mapper_fn.
 
-    Parameters:
+    Args:
         model: A PyTorch model.
         source_class: the layer class to replace. Each find will be passed to mapper_fn
         mapper_fn: A callable that takes as argument the layer to replace and returns the new object.

sinabs/from_torch.py

@@ -31,7 +31,7 @@ def from_model(
     analyzed, and a copy is returned, with all ReLUs and NeuromorphicReLUs turned into
     SpikingLayers.
 
-    Parameters:
+    Args:
         model: Torch model
         input_shape: If provided, the layer dimensions are computed. Otherwise they will be computed at the first forward pass.
         spike_threshold: The membrane potential threshold for spiking (same for all layers).
@@ -101,6 +101,28 @@ def from_model(
                     **kwargs_backend,
                 ).to(device),
             )
+
+        elif isinstance(model, nn.Module):
+            layers = [layer for _, layer in model.named_children()]
+
+            if not isinstance(layers[-1], (nn.ReLU, sl.NeuromorphicReLU)):
+                snn.add_module(
+                    "spike_output",
+                    spike_layer_class(
+                        spike_threshold=spike_threshold,
+                        spike_fn=spike_fn,
+                        reset_fn=reset_fn,
+                        surrogate_grad_fn=surrogate_grad_fn,
+                        min_v_mem=min_v_mem,
+                        **kwargs_backend,
+                    ).to(device),
+                )
+
+            else:
+                warn(
+                    "Spiking output can only be added to sequential models that do not end in a ReLU. No layer has been added."
+                )
+
         else:
             warn(
                 "Spiking output can only be added to sequential models that do not end in a ReLU. No layer has been added."

sinabs/hooks.py

@@ -12,7 +12,7 @@ from sinabs.layers import SqueezeMixin, StatefulLayer
 def _extract_single_input(input_data: List[Any]) -> Any:
     """Extract single element of a list.
 
-    Parameters:
+    Args:
         input_data: List that should have only one element
 
     Returns:
@@ -36,14 +36,13 @@ def conv_connection_map(
     """Generate connectivity map for a convolutional layer The map indicates for each element in
     the layer input to how many postsynaptic neurons it connects (i.e. the fanout)
 
-    Parameters:
-        layer: Convolutional layer for which connectivity map is to be
-               generated
+    Args:
+        layer: Convolutional layer for which connectivity map is to be generated
         input_shape: Shape of the input data (N, C, Y, X)
         output_shape: Shape of layer output given `input_shape`
         device: Device on which the connectivity map should reside.
-                Should be the same as that of the input to `layer`.
-                If None, will select device of the weight of `layer`.
+            Should be the same as that of the input to `layer`.
+            If None, will select device of the weight of `layer`.
 
     Returns:
         torch.Tensor: Connectivity map indicating the fanout for each
@@ -78,10 +77,9 @@ def get_hook_data_dict(module: nn.Module) -> Dict:
     """Convenience function to get `hook_data` attribute of a module if it has one and create it
     otherwise.
 
-    Parameters:
+    Args:
         module: The module whose `hook_data` dict is to be fetched.
-                If it does not have an attribute of that name, it
-                will add an empty dict.
+            If it does not have an attribute of that name, it will add an empty dict.
     Returns:
         The `hook_data` attribute of `module`. Should be a Dict.
     """
@@ -107,15 +105,15 @@ def input_diff_hook(
     at each forward pass. Afterwards the data can be accessed with
     `module.hook_data['diff_output']`
 
-    Parameters:
+    If `module` does not already have a `hook_data` attribute, it
+    will be added and the difference value described above  will be
+    stored under the key 'diff_output'. It is a tensor of the same
+    shape as `output`.
+
+    Args:
         module: Either a torch.nn.Conv2d or Linear layer
         input_: List of inputs to the layer. Should hold a single tensor.
         output: The layer's output.
-    Effect:
-        If `module` does not already have a `hook_data` attribute, it
-        will be added and the difference value described above  will be
-        stored under the key 'diff_output'. It is a tensor of the same
-        shape as `output`.
     """
     data = get_hook_data_dict(module)
     input_ = _extract_single_input(input_)
@@ -147,14 +145,14 @@ def firing_rate_hook(module: StatefulLayer, input_: Any, output: torch.Tensor):
     at each forward pass. Afterwards the data can be accessed with
     `module.hook_data['firing_rate']`
 
-    Parameters:
+    If `module` does not already have a `hook_data` attribute, it
+    will be added and the mean firing rate will be stored under the
+    key 'firing_rate'. It is a scalar value.
+
+    Args:
         module: A spiking sinabs layer, such as `IAF` or `LIF`.
         input_: List of inputs to the layer. Ignored here.
         output: The layer's output.
-    Effect:
-        If `module` does not already have a `hook_data` attribute, it
-        will be added and the mean firing rate will be stored under the
-        key 'firing_rate'. It is a scalar value.
     """
     data = get_hook_data_dict(module)
     data["firing_rate"] = output.mean()
@@ -171,16 +169,15 @@ def firing_rate_per_neuron_hook(
     `torch.register_forward_hook`. It will be called automatically
     at each forward pass. Afterwards the data can be accessed with
     `module.hook_data['firing_rate_per_neuron']`
+    If `module` does not already have a `hook_data` attribute, it
+    will be added and the mean firing rate will be stored under the
+    key 'firing_rate_per_neuron'. It is a tensor of the same
+    shape as neurons of the spiking layer.
 
-    Parameters:
+    Args:
         module: A spiking sinabs layer, such as `IAF` or `LIF`.
         input_: List of inputs to the layer. Ignored here.
         output: The layer's output.
-    Effect:
-        If `module` does not already have a `hook_data` attribute, it
-        will be added and the mean firing rate will be stored under the
-        key 'firing_rate_per_neuron'. It is a tensor of the same
-        shape as neurons of the spiking layer.
     """
     data = get_hook_data_dict(module)
     if isinstance(module, SqueezeMixin):
@@ -208,16 +205,16 @@ def conv_layer_synops_hook(
     at each forward pass. Afterwards the data can be accessed with
     `module.hook_data['layer_synops_per_timestep']`
 
-    Parameters:
+    If `module` does not already have a `hook_data` attribute, it
+    will be added and the mean firing rate will be stored under the
+    key 'layer_synops_per_timestep'. It is a scalar value.
+    It will also store a connectivity map under the key 'connection_map',
+    which holds the fanout for each input neuron.
+
+    Args:
         module: A torch.nn.Conv2d layer
         input_: List of inputs to the layer. Must contain exactly one tensor
         output: The layer's output.
-    Effect:
-        If `module` does not already have a `hook_data` attribute, it
-        will be added and the mean firing rate will be stored under the
-        key 'layer_synops_per_timestep'. It is a scalar value.
-        It will also store a connectivity map under the key 'connection_map',
-        which holds the fanout for each input neuron.
     """
     data = get_hook_data_dict(module)
     input_ = _extract_single_input(input_)
@@ -252,14 +249,14 @@ def linear_layer_synops_hook(
     at each forward pass. Afterwards the data can be accessed with
     `module.hook_data['layer_synops_per_timestep']`
 
-    Parameters:
+    If `module` does not already have a `hook_data` attribute, it
+    will be added and the mean firing rate will be stored under the
+    key 'layer_synops_per_timestep'.
+
+    Args:
         module: A torch.nn.Linear layer.
         input_: List of inputs to the layer. Must contain exactly one tensor
         output: The layer's output.
-    Effect:
-        If `module` does not already have a `hook_data` attribute, it
-        will be added and the mean firing rate will be stored under the
-        key 'layer_synops_per_timestep'.
     """
     data = get_hook_data_dict(module)
     input_ = _extract_single_input(input_)
@@ -308,7 +305,7 @@ class ModelSynopsHook:
       the synaptic operations per second, under the assumption that `dt`
       is the time step in seconds.
 
-    Parameters:
+    Attributes:
       dt: If not None, should be a float that indicates the simulation
           time step in seconds. The synaptic operations will be also
           provided in terms of synops per second.
@@ -320,7 +317,7 @@ class ModelSynopsHook:
         """Forward call of the synops model hook. Should not be called manually but only by PyTorch
         during a forward pass.
 
-        Parameters:
+        Args:
             module: A torch.nn.Sequential
             input_: List of inputs to the module.
             output: The module output.
@@ -380,7 +377,7 @@ def register_synops_hooks(module: nn.Sequential, dt: Optional[float] = None):
     This can be used instead of calling the torch function
     `register_forward_hook` on all layers.
 
-    Parameters:
+    Args:
         module: Sequential model for which the hooks should be registered.
         dt: If not None, should be a float indicating the simulation
             time step in seconds. Will also calculate synaptic operations per second.

sinabs/layers/alif.py

@@ -38,19 +38,19 @@ class ALIF(StatefulLayer):
         i(t+1) = \\alpha_{syn} i(t) (1-\\alpha_{syn}) + input
 
 
-    Parameters:
+    Args:
         tau_mem: Membrane potential time constant.
         tau_adapt: Spike threshold time constant.
         tau_syn: Synaptic decay time constants. If None, no synaptic dynamics are used, which is the default.
         adapt_scale: The amount that the spike threshold is bumped up for every spike, after which it decays back to the initial threshold.
         spike_threshold: Spikes are emitted if v_mem is above that threshold. By default set to 1.0.
         spike_fn: Choose a Sinabs or custom torch.autograd.Function that takes a dict of states,
-                  a spike threshold and a surrogate gradient function and returns spikes. Be aware
-                  that the class itself is passed here (because torch.autograd methods are static)
-                  rather than an object instance.
+            a spike threshold and a surrogate gradient function and returns spikes. Be aware
+            that the class itself is passed here (because torch.autograd methods are static)
+            rather than an object instance.
         reset_fn: A function that defines how the membrane potential is reset after a spike.
         surrogate_grad_fn: Choose how to define gradients for the spiking non-linearity during the
-                           backward pass. This is a function of membrane potential.
+            backward pass. This is a function of membrane potential.
         min_v_mem: Lower bound for membrane potential v_mem, clipped at every time step.
         shape: Optionally initialise the layer state with given shape. If None, will be inferred from input_size.
         train_alphas: When True, the discrete decay factor exp(-1/tau) is used for training rather than tau itself.
@@ -58,8 +58,8 @@ class ALIF(StatefulLayer):
         record_states: When True, will record all internal states such as v_mem or i_syn in a dictionary attribute `recordings`. Default is False.
 
     Shape:
-        - Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
-        - Output: Same as input.
+        Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
+        Output: Same as input.
 
     Attributes:
         v_mem: The membrane potential resets according to reset_fn for every spike.
@@ -137,7 +137,7 @@ class ALIF(StatefulLayer):
 
     def forward(self, input_data: torch.Tensor):
         """
-        Parameters:
+        Args:
             input_data: Data to be processed. Expected shape: (batch, time, ...)
 
         Returns:
@@ -244,7 +244,7 @@ class ALIFRecurrent(ALIF):
     .. math ::
         i(t+1) = \\alpha_{syn} i(t) (1-\\alpha_{syn}) + input
 
-    Parameters:
+    Args:
         tau_mem: Membrane potential time constant.
         tau_adapt: Spike threshold time constant.
         rec_connect: An nn.Module which defines the recurrent connectivity, e.g. nn.Linear
@@ -252,12 +252,12 @@ class ALIFRecurrent(ALIF):
         adapt_scale: The amount that the spike threshold is bumped up for every spike, after which it decays back to the initial threshold.
         spike_threshold: Spikes are emitted if v_mem is above that threshold. By default set to 1.0.
         spike_fn: Choose a Sinabs or custom torch.autograd.Function that takes a dict of states,
-                  a spike threshold and a surrogate gradient function and returns spikes. Be aware
-                  that the class itself is passed here (because torch.autograd methods are static)
-                  rather than an object instance.
+            a spike threshold and a surrogate gradient function and returns spikes. Be aware
+            that the class itself is passed here (because torch.autograd methods are static)
+            rather than an object instance.
         reset_fn: A function that defines how the membrane potential is reset after a spike.
         surrogate_grad_fn: Choose how to define gradients for the spiking non-linearity during the
-                           backward pass. This is a function of membrane potential.
+            backward pass. This is a function of membrane potential.
         min_v_mem: Lower bound for membrane potential v_mem, clipped at every time step.
         shape: Optionally initialise the layer state with given shape. If None, will be inferred from input_size.
         train_alphas: When True, the discrete decay factor exp(-1/tau) is used for training rather than tau itself.
@@ -265,8 +265,8 @@ class ALIFRecurrent(ALIF):
         record_states: When True, will record all internal states such as v_mem or i_syn in a dictionary attribute `recordings`. Default is False.
 
     Shape:
-        - Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
-        - Output: Same as input.
+        Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
+        Output: Same as input.
 
     Attributes:
         v_mem: The membrane potential resets according to reset_fn for every spike.
@@ -311,7 +311,7 @@ class ALIFRecurrent(ALIF):
 
     def forward(self, input_data: torch.Tensor):
         """
-        Parameters:
+        Args:
             input_data: Data to be processed. Expected shape: (batch, time, ...)
 
         Returns:

sinabs/layers/crop2d.py

@@ -9,7 +9,7 @@ ArrayLike = Union[np.ndarray, List, Tuple]
 class Cropping2dLayer(nn.Module):
     """Crop input image by.
 
-    Parameters:
+    Args:
         cropping: ((top, bottom), (left, right))
     """
 
@@ -38,7 +38,7 @@ class Cropping2dLayer(nn.Module):
     def get_output_shape(self, input_shape: Tuple) -> Tuple:
         """Retuns the output dimensions.
 
-        Parameters:
+        Args:
             input_shape: (channels, height, width)
 
         Returns:

sinabs/layers/exp_leak.py

@@ -17,7 +17,7 @@ class ExpLeak(LIF):
 
     where :math:`\\alpha =  e^{-1/tau_{mem}}` and :math:`\\sum z(t)` represents the sum of all input currents at time :math:`t`.
 
-    Parameters:
+    Args:
         tau_mem: Membrane potential time constant.
         min_v_mem: Lower bound for membrane potential v_mem, clipped at every time step.
         train_alphas: When True, the discrete decay factor exp(-1/tau) is used for training rather than tau itself.

sinabs/layers/iaf.py

@@ -22,24 +22,24 @@ class IAF(LIF):
 
     where :math:`\\sum z(t)` represents the sum of all input currents at time :math:`t`.
 
-    Parameters:
+    Args:
         spike_threshold: Spikes are emitted if v_mem is above that threshold. By default set to 1.0.
         spike_fn: Choose a Sinabs or custom torch.autograd.Function that takes a dict of states,
-                  a spike threshold and a surrogate gradient function and returns spikes. Be aware
-                  that the class itself is passed here (because torch.autograd methods are static)
-                  rather than an object instance.
+            a spike threshold and a surrogate gradient function and returns spikes. Be aware
+            that the class itself is passed here (because torch.autograd methods are static)
+            rather than an object instance.
         reset_fn: A function that defines how the membrane potential is reset after a spike.
         surrogate_grad_fn: Choose how to define gradients for the spiking non-linearity during the
-                           backward pass. This is a function of membrane potential.
+            backward pass. This is a function of membrane potential.
         tau_syn: Synaptic decay time constants. If None, no synaptic dynamics are used, which is the default.
         min_v_mem: Lower bound for membrane potential v_mem, clipped at every time step.
         shape: Optionally initialise the layer state with given shape. If None, will be inferred from input_size.
         record_states: When True, will record all internal states such as v_mem or i_syn in a dictionary attribute
-                       `recordings`. Default is False.
+            `recordings`. Default is False.
 
     Shape:
-        - Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
-        - Output: Same as input.
+        Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
+        Output: Same as input.
 
     Attributes:
         v_mem: The membrane potential resets according to reset_fn for every spike.
@@ -99,7 +99,7 @@ class IAFRecurrent(LIFRecurrent):
 
     where :math:`\\sum z(t)` represents the sum of all input currents at time :math:`t`.
 
-    Parameters:
+    Args:
         rec_connect: An nn.Module which defines the recurrent connectivity, e.g. nn.Linear
         spike_threshold: Spikes are emitted if v_mem is above that threshold. By default set to 1.0.
         spike_fn: Choose a Sinabs or custom torch.autograd.Function that takes a dict of states,
@@ -115,8 +115,8 @@ class IAFRecurrent(LIFRecurrent):
         record_states: When True, will record all internal states such as v_mem or i_syn in a dictionary attribute `recordings`. Default is False.
 
     Shape:
-        - Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
-        - Output: Same as input.
+        Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
+        Output: Same as input.
 
     Attributes:
         v_mem: The membrane potential resets according to reset_fn for every spike.

sinabs/layers/lif.py

@@ -31,7 +31,7 @@ class LIF(StatefulLayer):
     .. math ::
         \\text{if } V_{mem}(t) >= V_{th} \\text{, then } V_{mem} \\rightarrow V_{reset}
 
-    Parameters:
+    Args:
         tau_mem: Membrane potential time constant.
         tau_syn: Synaptic decay time constants. If None, no synaptic dynamics are used, which is the default.
         spike_threshold: Spikes are emitted if v_mem is above that threshold. By default set to 1.0.
@@ -50,8 +50,8 @@ class LIF(StatefulLayer):
             attribute `recordings`. Default is False.
 
     Shape:
-        - Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
-        - Output: Same as input.
+        Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
+        Output: Same as input.
 
     Attributes:
         v_mem: The membrane potential resets according to reset_fn for every spike.
@@ -155,7 +155,7 @@ class LIF(StatefulLayer):
 
     def forward(self, input_data: torch.Tensor) -> torch.Tensor:
         """
-        Parameters:
+        Args:
             input_data: Data to be processed. Expected shape: (batch, time, ...)
 
         Returns:
@@ -254,7 +254,7 @@ class LIFRecurrent(LIF):
     .. math ::
         \\text{if } V_{mem}(t) >= V_{th} \\text{, then } V_{mem} \\rightarrow V_{reset}
 
-    Parameters:
+    Args:
         tau_mem: Membrane potential time constant.
         rec_connect: An nn.Module which defines the recurrent connectivity, e.g. nn.Linear
         tau_syn: Synaptic decay time constants. If None, no synaptic dynamics are used, which is the default.
@@ -269,8 +269,8 @@ class LIFRecurrent(LIF):
         record_states: When True, will record all internal states such as v_mem or i_syn in a dictionary attribute `recordings`. Default is False.
 
     Shape:
-        - Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
-        - Output: Same as input.
+        Input: :math:`(Batch, Time, Channel, Height, Width)` or :math:`(Batch, Time, Channel)`
+        Output: Same as input.
 
     Attributes:
         v_mem: The membrane potential resets according to reset_fn for every spike.
@@ -309,8 +309,8 @@ class LIFRecurrent(LIF):
 
     def forward(self, input_data: torch.Tensor):
         """
-        Parameters:
-            input_data: Data to be processed. Expected shape: (batch, time, ...)
+        Args:
+            input_data: Data to be processed. Expected shape: (batch, time, ...).
 
         Returns:
             Output data with same shape as `input_data`.

sinabs/layers/neuromorphic_relu.py

@@ -7,14 +7,15 @@ class NeuromorphicReLU(torch.nn.Module):
     """NeuromorphicReLU layer. This layer is NOT used for Sinabs networks; it's useful while
     training analogue pyTorch networks for future use with Sinabs.
 
-    Parameters:
-        quantize: Whether or not to quantize the output (i.e. floor it to \
-        the integer below), in order to mimic spiking behavior.
-        fanout: Useful when computing the number of SynOps of a quantized \
-        NeuromorphicReLU. The activity can be accessed through \
-        NeuromorphicReLU.activity, and is multiplied by the value of fanout.
-        stochastic_rounding: Upon quantization, should the value be rounded stochastically or floored
-        Only done during training. During evaluation mode, the value is simply floored
+    Args:
+        quantize: Whether or not to quantize the output (i.e. floor it to the
+            integer below), in order to mimic spiking behavior.
+        fanout: Useful when computing the number of SynOps of a quantized
+            NeuromorphicReLU. The activity can be accessed through
+            NeuromorphicReLU.activity, and is multiplied by the value of fanout.
+        stochastic_rounding: Upon quantization, should the value be rounded
+            stochastically or floored. Only done during training. During
+            evaluation mode, the value is simply floored
     """
 
     def __init__(self, quantize=True, fanout=1, stochastic_rounding=False):

sinabs/layers/pool2d.py

@@ -74,7 +74,7 @@ class SpikingMaxPooling2dLayer(nn.Module):
     def get_output_shape(self, input_shape: Tuple) -> Tuple:
         """Returns the shape of output, given an input to this layer.
 
-        Parameters:
+        Args:
             input_shape: (channels, height, width)
 
         Returns:
@@ -95,10 +95,10 @@ class SumPool2d(torch.nn.LPPool2d):
     """Non-spiking sumpooling layer to be used in analogue Torch models. It is identical to
     torch.nn.LPPool2d with p=1.
 
-    Parameters:
-        kernel_size: the size of the window
-        stride: the stride of the window. Default value is kernel_size
-        ceil_mode: when True, will use ceil instead of floor to compute the output shape
+    Args:
+        kernel_size: the size of the window.
+        stride: the stride of the window. Default value is kernel_size.
+        ceil_mode: when True, will use ceil instead of floor to compute the output shape.
     """
 
     def __init__(self, kernel_size, stride=None, ceil_mode=False):

sinabs/layers/quantize.py

@@ -6,7 +6,7 @@ from sinabs.activation import Quantize
 class QuantizeLayer(nn.Module):
     """Layer that quantizes the input, i.e. returns floor(input).
 
-    Parameters:
+    Args:
         quantize: If False, this layer will pass on the input without modifying it.
     """
 

sinabs/layers/stateful_layer.py

@@ -8,7 +8,7 @@ class StatefulLayer(torch.nn.Module):
     """A base class that instantiates buffers/states which update at every time step and provides
     helper methods that manage those states.
 
-    Parameters:
+    Args:
         state_names: the PyTorch buffers to initialise. These are not parameters.
     """
 
@@ -97,12 +97,15 @@ class StatefulLayer(torch.nn.Module):
     ):
         """Reset the state/buffers in a layer.
 
-        Parameters:
-            randomize: If true, reset the states between a range provided. Else, the states are reset to zero.
-            value_ranges: A dictionary of key value pairs: buffer_name -> (min, max) for each state that needs to be reset.
-                          The states are reset with a uniform distribution between the min and max values specified.
-                          Any state with an undefined key in this dictionary will be reset between 0 and 1
-                          This parameter is only used if randomize is set to true.
+        Args:
+            randomize: If true, reset the states between a range provided.
+                Else, the states are reset to zero.
+            value_ranges: A dictionary of key value pairs: buffer_name -> (min,
+                max) for each state that needs to be reset. The states are
+                reset with a uniform distribution between the min and max
+                values specified. Any state with an undefined key in this
+                dictionary will be reset between 0 and 1. This parameter is
+                only used if randomize is set to true.
 
         .. note:: If you would like to reset the state with a custom distribution, you can do this individually for each parameter as follows::
 

sinabs/layers/to_spike.py

@@ -7,15 +7,15 @@ from torch import nn
 class Img2SpikeLayer(nn.Module):
     """Layer to convert images to spikes.
 
-    Parameters:
-        image_shape: tuple image shape
-        tw: int Time window length
-        max_rate: maximum firing rate of neurons
-        layer_name: string layer name
-        norm: the supposed maximum value of the input (default 255.0)
-        squeeze: whether to remove singleton dimensions from the input
+    Args:
+        image_shape: tuple image shape.
+        tw: int Time window length.
+        max_rate: maximum firing rate of neurons.
+        layer_name: string layer name.
+        norm: the supposed maximum value of the input (default 255.0).
+        squeeze: whether to remove singleton dimensions from the input.
         negative_spikes: whether to allow negative spikes in response
-                         to negative input
+            to negative input.
     """
 
     def __init__(
@@ -60,7 +60,7 @@ class Img2SpikeLayer(nn.Module):
 class Sig2SpikeLayer(torch.nn.Module):
     """Layer to convert analog Signals to spikes.
 
-    Parameters:
+    Args:
         channels_in: number of channels in the analog signal
         tw: int number of time steps for each sample of the signal (up sampling)
         layer_name: string layer name