Repository: jeshraghian/snntorch (language: python)

Issue:
Add power profiling capabilities
This seems to be a super popular feature request. 
Making accurate estimates seems near impossible, but we can probably generate an order of magnitude guess here.

The user would construct a model, pass data in, and the power profiling function returns the number of Synaptic operations in the forward-pass (this could be averaged across batches). 

Each synaptic op would be scaled by the energy cost for all selected devices; e.g., various GPUs & neuromorphic hardware. 
The same number would be given for non-spiking networks too. This could be achieved by just removing the spiking modules.

SpikingKeras has a similar function that does it really nicely. However, it overstates the improvement given with spikes because it does not account for overhead (i.e., moving data to/from memory, or between multiple chips).

Including an argument that factors in overhead would by tricky, but useful. 
The model would be parsed for number of neurons/synapses, and if either exceeds the bandwidth of a single chip, then we need to estimate how frequently data needs to be moved between chips & add that to the overall energy consumption.

A lot of coarse estimates would be made, but I think it could be helpful.
Add power profiling capabilities
This seems to be a super popular feature request. 
Making accurate estimates seems near impossible, but we can probably generate an order of magnitude guess here.

The user would construct a model, pass data in, and the power profiling function returns the number of Synaptic operations in the forward-pass (this could be averaged across batches). 

Each synaptic op would be scaled by the energy cost for all selected devices; e.g., various GPUs & neuromorphic hardware. 
The same number would be given for non-spiking networks too. This could be achieved by just removing the spiking modules.

SpikingKeras has a similar function that does it really nicely. However, it overstates the improvement given with spikes because it does not account for overhead (i.e., moving data to/from memory, or between multiple chips).

Including an argument that factors in overhead would by tricky, but useful. 
The model would be parsed for number of neurons/synapses, and if either exceeds the bandwidth of a single chip, then we need to estimate how frequently data needs to be moved between chips & add that to the overall energy consumption.

A lot of coarse estimates would be made, but I think it could be helpful.


Discussion / hints:
Dear Sir,

Are there plans or work being done to implement it? Can I implement it or help with the implementation ? 

Best Regards,
Dear Sir,

Are there plans or work being done to implement it? Can I implement it or help with the implementation ? 

Best Regards,

Fix patch:
diff --git a/snntorch/energy_estimation/FormattingOptions.py b/snntorch/energy_estimation/FormattingOptions.py
new file mode 100644
index 00000000..788a62a7
--- /dev/null
+++ b/snntorch/energy_estimation/FormattingOptions.py
@@ -0,0 +1,192 @@
+from __future__ import annotations
+from typing import Any, Dict, Iterable, Sequence, Union
+
+import math
+from typing import Any, List
+from .device_profile import DeviceProfile
+from .layer_info import LayerInfo
+import numpy as np
+import torch
+from torch import nn
+from torch.jit import ScriptModule
+from .enums import *
+
+"""
+    FormattingOptions didn't change much from torchinfo project https://github.com/TylerYep/torchinfo, all credits for
+    displaying very nicely architectures goes there
+"""
+
+try:
+    from torch.nn.parameter import is_lazy
+except ImportError:
+
+    def is_lazy(param: nn.Parameter) -> bool:  # type: ignore[misc]
+        del param
+        return False
+
+DETECTED_INPUT_OUTPUT_TYPES = Union[
+    Sequence[Any], Dict[Any, torch.Tensor], torch.Tensor
+]
+
+HEADER_TITLES = {
+    ColumnSettings.KERNEL_SIZE: "Kernel Shape",
+    ColumnSettings.INPUT_SIZE: "Input Shape",
+    ColumnSettings.OUTPUT_SIZE: "Output Shape",
+    ColumnSettings.NUMBER_OF_NEURONS: "Neurons #",
+    ColumnSettings.AVERAGE_FIRING_RATE: "Average Firing Rate[Hz]",
+    ColumnSettings.TOTAL_EVENTS: "Total events #",
+    ColumnSettings.SPIKING_EVENTS: "Spiking events #",
+    ColumnSettings.NUMBER_OF_SYNAPSES: "Synapses #",
+    ColumnSettings.NUM_PARAMS: "Param #",
+    ColumnSettings.PARAMS_PERCENT: "Param %",
+    ColumnSettings.MULT_ADDS: "Mult-Adds",
+    ColumnSettings.TRAINABLE: "Trainable",
+}
+CONVERSION_FACTORS = {
+    Units.TERABYTES: 1e12,
+    Units.GIGABYTES: 1e9,
+    Units.MEGABYTES: 1e6,
+    Units.NONE: 1,
+}
+
+
+class FormattingOptions:
+    """Class that holds information about formatting the table output."""
+
+    def __init__(
+        self,
+        max_depth: int,
+        verbose: int,
+        col_names: List[ColumnSettings, ...],
+        col_width: int,
+        row_settings: set[RowSettings],
+        devices: List[DeviceProfile]
+    ) -> None:
+        self.max_depth = max_depth
+        self.verbose = verbose
+        self.col_names = col_names
+        self.col_width = col_width
+        self.row_settings = row_settings
+        self.params_units = Units.NONE
+        self.macs_units = Units.AUTO
+        self.rounding_digits = 3
+
+        self.devices = devices
+        for device in self.devices:
+            HEADER_TITLES[device.name] = "energy for " + device.name + " [J]"
+            self.col_names.append(device.name)
+
+        self.layer_name_width = 40
+        self.ascii_only = RowSettings.ASCII_ONLY in self.row_settings
+        self.show_var_name = RowSettings.VAR_NAMES in self.row_settings
+        self.show_depth = RowSettings.DEPTH in self.row_settings
+        self.hide_recursive_layers = (
+            RowSettings.HIDE_RECURSIVE_LAYERS in self.row_settings
+        )
+
+
+    @staticmethod
+    def str_(val: Any) -> str:
+        return str(val) if val is not None else "--"
+
+    def get_start_str(self, depth: int) -> str:
+        """This function should handle all ascii/non-ascii-related characters."""
+        if depth == 0:
+            return ""
+        if depth == 1:
+            return "+ " if self.ascii_only else "├─"
+        return ("|    " if self.ascii_only else "│    ") * (depth - 1) + (
+            "+ " if self.ascii_only else "└─"
+        )
+
+    def set_layer_name_width(
+        self, summary_list: list[LayerInfo], align_val: int = 5
+    ) -> None:
+        """
+        Set layer name width by taking the longest line length and rounding up to
+        the nearest multiple of align_val.
+        """
+        max_length = 0
+        for info in summary_list:
+            depth_indent = info.depth * align_val + 1
+            layer_title = info.get_layer_name(self.show_var_name, self.show_depth)
+            max_length = max(max_length, len(layer_title) + depth_indent)
+        if max_length >= self.layer_name_width:
+            self.layer_name_width = math.ceil(max_length / align_val) * align_val
+
+    def get_total_width(self) -> int:
+        """Calculate the total width of all lines in the table."""
+        return len(tuple(self.col_names)) * self.col_width + self.layer_name_width
+
+    def format_row(self, layer_name: str, row_values: dict[ColumnSettings, str]) -> str:
+        """Get the string representation of a single layer of the model."""
+        info_to_use = [row_values.get(row_type, "") for row_type in self.col_names]
+        new_line = f"{layer_name:<{self.layer_name_width}} "
+        for info in info_to_use:
+            new_line += f"{info:<{self.col_width}} "
+
+        return new_line.rstrip() + "\n"
+
+    def header_row(self) -> str:
+        layer_header = ""
+        if self.show_var_name:
+            layer_header += " (var_name)"
+        if self.show_depth:
+            layer_header += ":depth-idx"
+
+        return self.format_row(f"Layer (type{layer_header})", HEADER_TITLES)
+
+    def layer_info_to_row(
+        self, layer_info: LayerInfo, reached_max_depth: bool, total_params: int
+    ) -> str:
+        """Convert layer_info to string representation of a row."""
+        values_for_row = {
+            ColumnSettings.KERNEL_SIZE: self.str_(layer_info.kernel_size),
+            ColumnSettings.INPUT_SIZE: self.str_(layer_info.input_size),
+            ColumnSettings.OUTPUT_SIZE: self.str_(layer_info.output_size),
+            ColumnSettings.NUMBER_OF_SYNAPSES: self.str_(layer_info.synapse_count),
+            ColumnSettings.NUMBER_OF_NEURONS: self.str_(layer_info.neuron_count),
+            ColumnSettings.SPIKING_EVENTS: self.str_(layer_info.spiking_events),
+            ColumnSettings.TOTAL_EVENTS: self.str_(layer_info.total_events),
+            ColumnSettings.AVERAGE_FIRING_RATE: self.str_(None if layer_info.firing_rate is None else
+                                                          round(layer_info.firing_rate, 4)),
+            ColumnSettings.NUM_PARAMS: layer_info.num_params_to_str(reached_max_depth),
+            ColumnSettings.PARAMS_PERCENT: layer_info.params_percent(
+                total_params, reached_max_depth
+            ),
+            ColumnSettings.MULT_ADDS: layer_info.macs_to_str(reached_max_depth),
+            ColumnSettings.TRAINABLE: self.str_(layer_info.trainable),
+        }
+
+        for i, device in enumerate(self.devices):
+            values_for_row[device.name] = self.str_(f"%.{self.rounding_digits}g" % layer_info.total_energy_contributions[i])
+
+        start_str = self.get_start_str(layer_info.depth)
+        layer_name = layer_info.get_layer_name(self.show_var_name, self.show_depth)
+        new_line = self.format_row(f"{start_str}{layer_name}", values_for_row)
+
+        if self.verbose == Verbosity.VERBOSE:
+            for inner_name, inner_layer_info in layer_info.inner_layers.items():
+                prefix = self.get_start_str(layer_info.depth + 1)
+                new_line += self.format_row(f"{prefix}{inner_name}", inner_layer_info)
+        return new_line
+
+    def layers_to_str(self, summary_list: list[LayerInfo], total_params: int) -> str:
+        """
+        Print each layer of the model using only current layer info.
+        Container modules are already dealt with in add_missing_container_layers.
+        """
+        new_str = ""
+        for layer_info in summary_list:
+            if (
+                layer_info.depth > self.max_depth
+                or self.hide_recursive_layers
+                and layer_info.is_recursive
+            ):
+                continue
+
+            reached_max_depth = layer_info.depth == self.max_depth
+            new_str += self.layer_info_to_row(
+                layer_info, reached_max_depth, total_params
+            )
+        return new_str
diff --git a/snntorch/energy_estimation/__init__.py b/snntorch/energy_estimation/__init__.py
new file mode 100644
index 00000000..aaa32852
--- /dev/null
+++ b/snntorch/energy_estimation/__init__.py
@@ -0,0 +1,1 @@
+import torch
diff --git a/snntorch/energy_estimation/device_profile.py b/snntorch/energy_estimation/device_profile.py
new file mode 100644
index 00000000..31e49fcb
--- /dev/null
+++ b/snntorch/energy_estimation/device_profile.py
@@ -0,0 +1,28 @@
+class DeviceProfile(object):
+    """
+        A class representing the information about the device, which is used when calculating the energy estimates.
+        It considers few properties of the device :
+
+        :param name : unique identifier of the device, used in estimate_energy
+
+        :param is_spiking_architecture : is the device neuromorphic. If set to true, energy is added only when there is
+        an event (synaptic or neuron). for example, for final output layer with [1, 0, 0, 0, 0, 0] there is only one
+        event.
+
+        :param energy_per_synapse_event: energy required to perform a synapse event. This is when a input parameter
+        is multiplied by certain weight in Linear or Convolutional layers. Units are J / event
+
+        :param energy_per_neuron_event : energy required to perform a neuron event. This is energy required, when an
+        output is generated from a neuron
+
+    """
+
+    def __init__(self, name : str, energy_per_synapse_event: float, energy_per_neuron_event: float, is_spiking_architecture: bool):
+        self.name = name
+        self.energy_per_synapse_event = energy_per_synapse_event
+        self.energy_per_neuron_event = energy_per_neuron_event
+        self.is_spiking_architecture = is_spiking_architecture
+
+    def __str__(self):
+        return f"name={self.name}, synapse_event_energy={self.energy_per_synapse_event}, " \
+               f"neuron_event_energy={self.energy_per_neuron_event}, isSpiking={self.is_spiking_architecture}"
diff --git a/snntorch/energy_estimation/device_profile_registry.py b/snntorch/energy_estimation/device_profile_registry.py
new file mode 100644
index 00000000..160a3d45
--- /dev/null
+++ b/snntorch/energy_estimation/device_profile_registry.py
@@ -0,0 +1,50 @@
+from .device_profile import DeviceProfile
+from typing import Dict, List
+import logging
+
+
+class DeviceWithNameDoesntExist(Exception):
+    """Raised when the get_device was called with name that doesn't map to any DeviceProfile"""
+
+
+class DeviceProfileRegistry:
+    """
+        A singleton with a dictionary _devices_reg consisting of all registered/known devices (a map from name to
+        DeviceProfile).
+
+        Example::
+
+            # add the device to Registry
+            DeviceProfileRegistry.add_device('new-cpu', 1e-9, 1e-9, False)
+
+            # Specify your network
+            # ...
+
+            # show the estimates for that model
+            estimate_energy(model=my_pytorch_model, devices="new-cpu")
+
+    """
+
+    _devices_reg: Dict[str, DeviceProfile] = {"cpu": DeviceProfile("cpu", 8.6e-9, 8.6e-9, False)
+                                              }
+
+    @staticmethod
+    def get_device(name) -> DeviceProfile:
+        if name not in DeviceProfileRegistry._devices_reg:
+            raise DeviceWithNameDoesntExist()
+
+        return DeviceProfileRegistry._devices_reg[name]
+
+    @staticmethod
+    def add_device(name: str, energy_per_synapse_event: float, energy_per_neuron_event: float,
+                   is_spiking_architecture: bool) -> None:
+
+        if name in DeviceProfileRegistry._devices_reg:
+            logging.warning(f"Device with name {name} already exists, overwriting ...")
+
+        DeviceProfileRegistry._devices_reg[name] = DeviceProfile(name, energy_per_synapse_event,
+                                                                 energy_per_neuron_event, is_spiking_architecture)
+
+    @staticmethod
+    def get_all_registered_devices() -> List[str]:
+        return list(DeviceProfileRegistry._devices_reg.keys())
diff --git a/snntorch/energy_estimation/energy_estimation_network_interface.py b/snntorch/energy_estimation/energy_estimation_network_interface.py
new file mode 100644
index 00000000..541e9a1d
--- /dev/null
+++ b/snntorch/energy_estimation/energy_estimation_network_interface.py
@@ -0,0 +1,22 @@
+import torch
+from torch.nn import Module
+
+
+class EnergyEstimationNetworkInterface(Module):
+    """
+        An interface which makes the estimate_energy easier, because it will use the .reset function which should
+        reset the potentials on the membrane (or in general hidden states), therefore guaranteeing consistent results
+        when doing the forward passes on given spike train (before doing forward passes, estimate_energy will reset
+        the network).
+
+        TODO : this is not a great solution but it currently works
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self):
+        raise NotImplemented()
+
+    def reset(self):
+        raise NotImplemented()
diff --git a/snntorch/energy_estimation/enums.py b/snntorch/energy_estimation/enums.py
new file mode 100644
index 00000000..72d709b1
--- /dev/null
+++ b/snntorch/energy_estimation/enums.py
@@ -0,0 +1,48 @@
+from enum import Enum, IntEnum, unique
+
+"""
+    Enums below has been mostly taken from torchinfo project https://github.com/TylerYep/torchinfo, all credits for
+    displaying very nicely architectures goes there
+"""
+
+@unique
+class Units(str, Enum):
+    """Enum containing all available bytes units."""
+
+    AUTO = "auto"
+    MEGABYTES = "M"
+    GIGABYTES = "G"
+    TERABYTES = "T"
+    NONE = ""
+
+@unique
+class Verbosity(IntEnum):
+    """Contains verbosity levels."""
+    QUIET, DEFAULT, VERBOSE = 0, 1, 2
+
+@unique
+class RowSettings(str, Enum):
+    """Enum containing all available row settings."""
+
+    DEPTH = "depth"
+    VAR_NAMES = "var_names"
+    ASCII_ONLY = "ascii_only"
+    HIDE_RECURSIVE_LAYERS = "hide_recursive_layers"
+
+
+@unique
+class ColumnSettings(str, Enum):
+    """Enum containing all available column settings."""
+
+    KERNEL_SIZE = "kernel_size"
+    INPUT_SIZE = "input_size"
+    OUTPUT_SIZE = "output_size"
+    NUM_PARAMS = "num_params"
+    PARAMS_PERCENT = "params_percent"
+    MULT_ADDS = "mult_adds"
+    TRAINABLE = "trainable"
+    NUMBER_OF_NEURONS = "neurons"
+    NUMBER_OF_SYNAPSES = "synapses"
+    AVERAGE_FIRING_RATE = "firing_rate"
+    SPIKING_EVENTS = "spiking_events"
+    TOTAL_EVENTS = "total_events"
diff --git a/snntorch/energy_estimation/estimate_energy.py b/snntorch/energy_estimation/estimate_energy.py
new file mode 100644
index 00000000..79bf8b26
--- /dev/null
+++ b/snntorch/energy_estimation/estimate_energy.py
@@ -0,0 +1,466 @@
+import numpy as np
+import torch
+import torch.nn as nn
+from typing import Sequence, Union, Iterable, Callable, Mapping, cast, List
+from .layer_info import LayerInfo, get_children_layers
+from .FormattingOptions import *
+from .model_statistics import *
+from .device_profile_registry import DeviceProfileRegistry
+from .energy_estimation_network_interface import EnergyEstimationNetworkInterface
+from snntorch.utils import reset
+from .utils import *
+import sys
+
+"""
+    The code from torchinfo project (https://github.com/TylerYep/torchinfo) has been adjusted to get the estimate energy
+    function (quite a bit of it was modified, however overall structure, especially with displaying is the same)
+"""
+
+
+def estimate_energy(model: Union[nn.Module , EnergyEstimationNetworkInterface],
+                    devices: Union[DeviceProfile , List[Union[DeviceProfile , str]] , str , None] = "cpu",
+                    input_size: Union[Sequence[int] , None] = None,
+                    input_data: Union[Union[torch.Tensor , np.ndarray , Sequence[float]] , None] = None,
+                    input_size_proportion_of_one: float = 1.0,
+                    verbose: Verbosity = Verbosity.VERBOSE,
+                    depth=1,
+                    col_width: int = 25,
+                    first_dim_is_time=True,
+                    network_requires_first_dim_as_time=True,
+                    include_bias_term_in_events=True,
+                    columns: List[ColumnSettings] = (ColumnSettings.NUMBER_OF_SYNAPSES,
+                                                     ColumnSettings.NUMBER_OF_NEURONS,
+                                                     ColumnSettings.AVERAGE_FIRING_RATE,
+                                                     ColumnSettings.TOTAL_EVENTS, ColumnSettings.SPIKING_EVENTS)):
+    """
+`       The main function handling generating the energy estimation feature. It takes a model, prepared by user input
+        (either specify a shape by passing `input_size` or data `input_data` on which we want the network to run ), and
+        a list of devices (represented by names or Device Profiles) with information how the energy should be calculated.
+
+        It is expected that the network either accepts input in shape (time, batchsize, ...) when
+        `network_requires_first_dim_as_time` is True or (batchsize, ...) when False.
+
+        :param model : a PyTorch architecture for which we want to do the energy estimation. Because the network itself
+        can be implemented in many ways (more flexibility than Keras), a helpful assumption is that the network except
+        forward method (or overall some method doing computation) it will also have a `reset` method from
+        `EnergyEstimationNetworkInterface` which will reset all hidden states from the network. If the module is
+        an instance of nn.Module, it will try to reset it using snntorch.utils reset function (TODO : however currently
+        I'm probably using it wrong, as I didn't get consistent numbers)
+
+        :param devices : a string or list of strings representing a name of the device (from DeviceProfileRegistry).
+        For every specified device, layer-wise energy estimation will be provided
+
+        :param input_size : a tuple of ints representing the shape of sample input. Setting this parameters to something
+        else than None is mutually exclusive with `input_data` (only one should be specified), please specify only one.
+        A tensor with requested shape will be prepared to pass to the network (each element of that tensor will have
+        `input_size_proportion_of_one`probability of being 1.0 (otherwise zero)).
+
+        :param input_data :a pytorch tensor (or python array of floats or numpy float array) that will be passed to the
+        Pytorch model `model`. Mutually exclusive with input_size, please specify only one
+
+        :param input_size_proportion_of_one : if `input_size` is provided, each of the elements will have probability
+        of being one set to this value.  it is implemented in following way :
+
+            >>> prob_mat = input_size_proportion_of_one * torch.ones(input_size)
+            >>> return torch.bernoulli(prob_mat)
+
+        :param first_dim_is_time : boolean flag which when set to true, will let this function know that the first
+        dimension in input (either input_data or input_shape) we pass to network is time dimension. It is important
+        to set it appropriately, so this function can interpret the data correctly
+
+        :param network_requires_first_dim_as_time : boolean flag which when set to true, will let this function know
+        that the PyTorch model requires first dimension to be time ( accepts (time, batchsize, ...)). If set to True,
+        an input in format (batchsize, ...) will be passed. It is important to set it appropriately, so the input data
+        can be passed correctly to PyTorch model.
+
+        :param include_bias_term_in_events : If set to False, the bias calculations for energy are excluded (makes them
+        the same as in kerasSpiking ). If set to True, bias terms are included
+
+        :param columns : a list of Columns to display
+
+        :param col_width : width of the column required for printing
+
+        :param depth : The level of Depth through which we will go to print the model.
+
+        :param verbose : verbosity level set for formatting options
+    """
+
+    # to make forward pass, we need to have either input data or input size
+    # assert input_data is not None or input_size is not None, "To run estimate energy"
+    validate_user_parameters(model, input_size, input_data, input_size_proportion_of_one)
+
+    # based on provided data, prepare the input for passing to the network
+    # we want our input to be in format (time, batch size , ... )
+    input_data = create_input_for_network_from_size(input_data, input_size, input_size_proportion_of_one,
+                                                    first_dim_is_time)
+
+    # get the information about the devices for which we will get the energy estimates
+    devices = _resolve_devices(devices)
+
+    # get the layer-wise information about the parameters/events/energies by doing a forward pass with the prepared
+    # `input_data` on PyTorch model `model`
+    summary_list = forward_pass(model, input_data, devices, first_dim_is_time,
+                                network_requires_first_dim_as_time, include_bias_term_in_events)
+
+    # displaying options from https://github.com/TylerYep/torchinfo
+    columns = list(columns)
+    formatting = FormattingOptions(depth, verbose, columns, col_width, {RowSettings.DEPTH}, devices)
+    results = ModelStatistics(summary_list, input_data.size(), get_total_memory_used(input_data), formatting,
+                              devices)
+    return results
+
+
+def validate_user_parameters(model: Union[nn.Module , EnergyEstimationNetworkInterface],
+                             input_size: Union[Sequence[int] , None] = None,
+                             input_data: Union[torch.Tensor, np.ndarray, Sequence[float], None] = None,
+                             input_size_proportion_of_one: float = 1.0):
+    """
+        Function which takes information about the data (either input_size or input_data) and does some basic sanity
+        checks about the data
+
+        TODO: probably add more here, inform user about potential problems
+    """
+    if input_size is not None and input_data is not None:
+        raise Exception("only one of input_size and input_data should be specified")
+
+    if input_size is None and input_data is None:
+        raise Exception(
+            "both input_size and input_data are None ! please specify one in order to do a forward pass and get "
+            "energy estimates")
+
+
+def create_input_for_network_from_size(input_data: Union[torch.Tensor , np.ndarray , Sequence[float], None] = None,
+                                       input_size: Union[Sequence[int], None] = None,
+                                       input_size_proportion_of_one: float = 1.0,
+                                       first_dim_is_time: bool = True) -> torch.Tensor:
+    """
+        Based on provided input, return a correct input to our network, in format (time, batchsize, ....)
+    """
+
+    if input_data is None:
+        if first_dim_is_time:
+            prob_mat = input_size_proportion_of_one * torch.ones(input_size)
+        else:
+            prob_mat = input_size_proportion_of_one * torch.ones((1, *input_size))
+        return torch.bernoulli(prob_mat)
+    else:
+        if first_dim_is_time:
+            return input_data
+        else:
+            return torch.stack([input_data])
+
+
+def forward_pass(model: Union[nn.Module , EnergyEstimationNetworkInterface], fwd_inp: torch.Tensor,
+                 devices: List[DeviceProfile], first_dim_is_time: bool, network_requires_first_dim_as_time: bool,
+                 include_bias_term_in_events: bool):
+    # get the class name of the model
+    model_name = model.__class__.__name__
+
+    # apply the forward hooks and pre-hooks
+    summary_list, global_layer_info, hooks = apply_hooks(model_name, model, fwd_inp, include_bias_term_in_events,
+                                                         network_requires_first_dim_as_time)
+
+    # set the model to evaluation model (no need for gradients)
+    model.eval()
+
+    # until this point the hooks has been registered, and input has been preprocessed to be in expected
+    # shape of (time, batch size, ...). If the flag `network_requires_first_dim_as_time` is true we will pass
+    # (time, batch size, ...) or when false just (batch size, ...)
+    if network_requires_first_dim_as_time:
+        # pass (timestep, batch size, ...)
+        fwd_inp_single = torch.stack([fwd_inp[0]])
+        with torch.no_grad():
+            model(fwd_inp_single)
+    else:
+        # pass (timestep, batch size, ...)
+        with torch.no_grad():
+            model(fwd_inp[0])
+
+    # set the children layer information to any nested layers
+    set_children_layers(summary_list)
+
+    # set the synapse and neuron counts for any nested layers
+    set_the_synapses_neurons_for_nested_layers(summary_list)
+
+    # remove the pre-hooks, they are no longer needed
+    for module_id, h in hooks.items():
+        h[0].remove()
+
+    # flip the settings for layers to estimating the energy
+    for module_id, layer_info in global_layer_info.items():
+        layer_info.calculate_energy_mode = True
+
+        # set the information about the devices to calculate the contributing energy
+        layer_info.setup_device_profiles(devices)
+
+    # TODO : this probably needs to be done better
+    if isinstance(model, EnergyEstimationNetworkInterface):
+        model.reset()
+    else:
+        # TODO : does this work as expected ? weird results ?
+        reset(model)
+
+    # get time number of timesteps
+    T = fwd_inp.size()[0]
+
+    # iterate over all timesteps
+
+    # TODO: should I pass the whole input here if the network requires time as first dimension ?
+    # something like this :
+    # if network_requires_first_dim_as_time:
+    #   model(fwd_inp)
+    # rather than below
+    for t in range(T):
+        if network_requires_first_dim_as_time:
+            inp = torch.stack([fwd_inp[t]])
+        else:
+            inp = fwd_inp[t]
+        model(inp)
+
+    # the forward passes for whole input has been made, calculate the energies
+    calculate_total_energies_nested_layers(summary_list)
+
+    return summary_list
+
+
+def set_the_synapses_neurons_for_nested_layers(summary_list: List[LayerInfo]):
+    # set the synapses correctly
+    for layer in summary_list:
+
+        # if neuron count is None (so either a non-basic module or unrecognized module), go through it's children
+        # and get the total number of neurons (and set its own neuron count to that value)
+        if layer.neuron_count is None:
+            n = None
+            for child in layer.children:
+                if child.neuron_count is not None:
+                    if n is None:
+                        n = 0
+                    n += child.neuron_count
+            layer.neuron_count = n
+
+        # do the same calculation for synapse count
+        if layer.synapse_count is None:
+            s = None
+            for child in layer.children:
+                if child.synapse_count is not None:
+                    if s is None:
+                        s = 0
+                    s += child.synapse_count
+            layer.synapse_count = s
+
+
+def calculate_total_energies_nested_layers(summary_list: List[LayerInfo]):
+    """
+        Calculate the total energy, spiking and total events for nested layers
+    """
+
+    # always add the energy contributions for the
+    for layer in summary_list:
+
+        # convert energy to numpy array (easier to add everything together)
+        total_energy = np.array(layer.total_energy_contributions, dtype=float)
+
+        # add the energies of its children to the total energy
+        for child in layer.children:
+            total_energy += np.array(child.total_energy_contributions, dtype=float)
+
+        # update the total energy contributions
+        layer.total_energy_contributions = total_energy
+
+        if layer.total_events is None:
+            total_events = 0
+            for child in layer.children:
+
+                # watch out for layers/children that were not recognized (i.e. the total events is none)
+                if child.total_events is not None:
+                    total_events += child.total_events
+            layer.total_events = None if total_events == 0 else total_events
+
+        if layer.spiking_events is None:
+            spiking_events = None
+            for child in layer.children:
+
+                # watch out for layers/children that were not recognized (i.e. the spiking events is none)
+                if child.spiking_events is not None:
+                    if spiking_events is None:
+                        spiking_events = 0
+                    spiking_events += child.spiking_events
+
+            layer.spiking_events = None if spiking_events is None else spiking_events
+
+        # set the firing rate to the spiking_events / total_events
+        if layer.spiking_events is not None and layer.total_events is not None and layer.firing_rate is None:
+            layer.firing_rate = layer.spiking_events / layer.total_events
+
+
+def construct_pre_hook(global_layer_info,
+                       summary_list,
+                       layer_ids,
+                       var_name,
+                       curr_depth,
+                       parent_info,
+                       include_bias_term_in_events,
+                       network_requires_first_dim_as_time):
+    def pre_hook(module: nn.Module, inputs) -> None:
+        """Create a Layer info object to aggregate layer information"""
+        del inputs
+
+        info = LayerInfo(var_name, module, curr_depth, parent_info,
+                         include_bias_term_in_events,
+                         network_requires_first_dim_as_time)
+        info.calculate_num_params()
+        info.check_recursive(layer_ids)
+        summary_list.append(info)
+        layer_ids.add(info.layer_id)
+        global_layer_info[info.layer_id] = info
+
+    return pre_hook
+
+
+def construct_hook(global_layer_info, batch_dim):
+    def hook(module: nn.Module, inputs, outputs) -> None:
+        """Update LayerInfo after forward pass."""
+        info = global_layer_info[id(module)]
+        if not info.calculate_energy_mode:
+            if info.contains_lazy_param:
+                info.calculate_num_params()
+
+            info.calculate_synapse_neuron_count(inputs, outputs)
+            info.input_size, _ = info.calculate_size(inputs, batch_dim)
+            info.output_size, elem_bytes = info.calculate_size(outputs, batch_dim)
+            info.output_bytes = elem_bytes * prod(info.output_size)
+            info.executed = True
+            # TODO: doesn't work
+            info.calculate_macs()
+        else:
+            info.calculate_events(inputs, outputs)
+
+    return hook
+
+
+def apply_hooks(model_name: str, model: nn.Module, fwd_inp: torch.Tensor, include_bias_term_in_events: bool,
+                network_requires_first_dim_as_time: bool):
+    """
+        recursively adds hooks to all layers of the model
+    """
+    summary_list = []
+    layers_ids = set()
+    global_layer_info = {}
+    hooks = {}
+
+    #        (modelname, nn.Module, depth, parent)
+    stack = [(model_name, model, 0, None)]
+
+    while stack:
+        module_name, layer, curr_depth, parent_info = stack.pop()
+        module_id = id(layer)
+
+        global_layer_info[module_id] = LayerInfo(module_name, layer, curr_depth, parent_info)
+
+        pre_hook = construct_pre_hook(global_layer_info, summary_list, layers_ids, module_name, curr_depth,
+                                      parent_info, include_bias_term_in_events=include_bias_term_in_events,
+                                      network_requires_first_dim_as_time=network_requires_first_dim_as_time
+                                      )
+
+        # register the hook using the last layer that uses this module
+        if module_id in hooks:
+            for hook in hooks[module_id]:
+                hook.remove()
+
+        hooks[module_id] = (layer.register_forward_pre_hook(pre_hook),
+                            layer.register_forward_hook(construct_hook(global_layer_info, batch_dim=0)),)
+
+        for name, mod in reversed(layer._modules.items()):
+            if mod is not None:
+                stack += [(name, mod, curr_depth + 1, global_layer_info[module_id])]
+
+    return summary_list, global_layer_info, hooks
+
+
+def traverse_input_data(
+    data: Any, action_fn: Callable[..., Any], aggregate_fn: Callable[..., Any]
+) -> Any:
+    """
+    Traverses any type of nested input data. On a tensor, returns the action given by
+    action_fn, and afterward aggregates the results using aggregate_fn.
+    """
+    if isinstance(data, torch.Tensor):
+        result = action_fn(data)
+    elif isinstance(data, np.ndarray):
+        result = action_fn(torch.from_numpy(data))
+        # If the result of action_fn is a torch.Tensor, then action_fn was meant for
+        #   torch.Tensors only (like calling .to(...)) -> Ignore.
+        if isinstance(result, torch.Tensor):
+            result = data
+
+    # Recursively apply to collection items
+    elif isinstance(data, Mapping):
+        aggregate = aggregate_fn(data)
+        result = aggregate(
+            {
+                k: traverse_input_data(v, action_fn, aggregate_fn)
+                for k, v in data.items()
+            }
+        )
+    elif isinstance(data, tuple) and hasattr(data, "_fields"):  # Named tuple
+        aggregate = aggregate_fn(data)
+        result = aggregate(
+            *(traverse_input_data(d, action_fn, aggregate_fn) for d in data)
+        )
+    elif isinstance(data, Iterable) and not isinstance(data, str):
+        aggregate = aggregate_fn(data)
+        result = aggregate(
+            [traverse_input_data(d, action_fn, aggregate_fn) for d in data]
+        )
+    else:
+        # Data is neither a tensor nor a collection
+        result = data
+    return result
+
+
+def get_total_memory_used(data) -> int:
+    """Calculates the total memory of all tensors stored in data."""
+    result = traverse_input_data(
+        data,
+        action_fn=lambda data: sys.getsizeof(
+            data.untyped_storage()
+            if hasattr(data, "untyped_storage")
+            else data.storage()
+        ),
+        aggregate_fn=(
+            # We don't need the dictionary keys in this case
+            lambda data: (lambda d: sum(d.values()))
+            if isinstance(data, Mapping)
+            else sum
+        ),
+    )
+    return cast(int, result)
+
+
+def _resolve_devices_base_case(device: Union[str , DeviceProfile]):
+    if isinstance(device, DeviceProfile):
+        return device
+    elif isinstance(device, str):
+        return DeviceProfileRegistry.get_device(device.lower())
+
+
+def _resolve_devices(devices: Union[str , DeviceProfile , List[Union[str , DeviceProfile]]]) -> List[DeviceProfile]:
+    if devices is None or len(devices) == 0:
+        return []
+
+    if isinstance(devices, str) or isinstance(devices, DeviceProfile):
+        return [_resolve_devices_base_case(devices)]
+
+    resolved_devices = []
+    for device in devices:
+        resolved_devices.append(_resolve_devices_base_case(device))
+    return resolved_devices
+
+
+def set_children_layers(summary_list: List[LayerInfo]) -> None:
+    """Populates the children and depth_index fields of all LayerInfo."""
+    idx: dict[int, int] = {}
+    for i, layer in enumerate(summary_list):
+        idx[layer.depth] = idx.get(layer.depth, 0) + 1
+        layer.depth_index = idx[layer.depth]
+        layer.children = get_children_layers(summary_list, i)
diff --git a/snntorch/energy_estimation/layer_info.py b/snntorch/energy_estimation/layer_info.py
new file mode 100644
index 00000000..7c7fa83d
--- /dev/null
+++ b/snntorch/energy_estimation/layer_info.py
@@ -0,0 +1,482 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, Sequence, Union, List
+from .device_profile import DeviceProfile
+
+import numpy as np
+import snntorch
+import torch
+from torch import nn
+from torch.jit import ScriptModule
+from snntorch.energy_estimation.layer_parameter_event_calculator import LayerParameterEventCalculator
+import logging
+
+from .enums import ColumnSettings
+
+
+"""
+    Layer info has been adjusted from torchinfo project https://github.com/TylerYep/torchinfo
+"""
+
+try:
+    from torch.nn.parameter import is_lazy
+except ImportError:
+
+    def is_lazy(param: nn.Parameter) -> bool:  # type: ignore[misc]
+        del param
+        return False
+
+DETECTED_INPUT_OUTPUT_TYPES = Union[
+    Sequence[Any], Dict[Any, torch.Tensor], torch.Tensor
+]
+
+
+class LayerInfo:
+    """Class that holds information about a layer module."""
+
+    def __init__(
+        self,
+        var_name: str,
+        module: nn.Module,
+        depth: int,
+        parent_info: Union[LayerInfo, None] = None,
+        include_bias_in_events: bool = True,
+        network_requires_first_dim_as_time : bool = True
+    ) -> None:
+        # Identifying information
+        self.layer_id = id(module)
+        self.module = module
+        self.class_name = (
+            str(module.original_name)
+            if isinstance(module, ScriptModule)
+            else module.__class__.__name__
+        )
+        # {layer name: {col_name: value_for_row}}
+        self.inner_layers: dict[str, dict[ColumnSettings, Any]] = {}
+        self.depth = depth
+        self.depth_index: Union[int , None] = None  # set at the very end
+        self.children: list[LayerInfo] = []  # set at the very end
+        self.executed = False
+        self.parent_info = parent_info
+        self.var_name = var_name
+        self.is_leaf_layer = not any(self.module.children())
+        self.contains_lazy_param = False
+        self.device_profiles: List[DeviceProfile] = []
+        self.total_energy_contributions: List[float] = []
+        self.include_bias_in_events = include_bias_in_events
+        self.network_requires_first_dim_as_time = network_requires_first_dim_as_time
+
+        # Statistics
+        self.is_recursive = False
+        self.input_size: list[int] = []
+        self.output_size: list[int] = []
+        self.kernel_size = self.get_kernel_size(module)
+        self.trainable_params = 0
+        self.num_params = 0
+        self.param_bytes = 0
+
+        self.spiking_events = None
+        self.total_events = None
+
+        self.output_bytes = 0
+        self.macs = 0
+        self.firing_rate = None
+        self.timestep = 0
+
+        self.neuron_count = None
+        self.synapse_count = None
+        self.calculate_energy_mode = False
+        self.is_synaptic = True
+
+    def __repr__(self) -> str:
+        return f"{self.class_name}: {self.depth}"
+
+    def setup_device_profiles(self, device_profiles: List[DeviceProfile]):
+        self.device_profiles = device_profiles
+        self.total_energy_contributions = [0] * len(device_profiles)
+
+    @property
+    def trainable(self) -> str:
+        """
+        Checks if the module is trainable. Returns:
+            "True", if all the parameters are trainable (`requires_grad=True`)
+            "False" if none of the parameters are trainable.
+            "Partial" if some weights are trainable, but not all.
+            "--" if no module has no parameters, like Dropout.
+        """
+        if self.num_params == 0:
+            return "--"
+        if self.trainable_params == 0:
+            return "False"
+        if self.num_params == self.trainable_params:
+            return "True"
+        if self.num_params > self.trainable_params:
+            return "Partial"
+        raise RuntimeError("Unreachable trainable calculation.")
+
+    def set_calculate_energy_mode(self):
+        self.calculate_energy_mode = True
+
+    @staticmethod
+    def calculate_size(
+        inputs: DETECTED_INPUT_OUTPUT_TYPES, batch_dim: Union[int , None]
+    ) -> tuple[list[int], int]:
+        """
+        Set input_size or output_size using the model's inputs.
+        Returns the corrected shape of `inputs` and the size of
+        a single element in bytes.
+        """
+        if inputs is None:
+            size, elem_bytes = [], 0
+
+        # pack_padded_seq and pad_packed_seq store feature into data attribute
+        elif (
+            isinstance(inputs, (list, tuple))
+            and inputs
+            and hasattr(inputs[0], "data")
+            and hasattr(inputs[0].data, "size")
+        ):
+            size = list(inputs[0].data.size())
+            elem_bytes = inputs[0].data.element_size()
+            if batch_dim is not None:
+                size = size[:batch_dim] + [1] + size[batch_dim + 1:]
+
+        elif isinstance(inputs, dict):
+            output = list(inputs.values())[-1]
+            size, elem_bytes = nested_list_size(output)
+            if batch_dim is not None:
+                size = [size[:batch_dim] + [1] + size[batch_dim + 1:]]
+
+        elif isinstance(inputs, torch.Tensor):
+            size = list(inputs.size())
+            elem_bytes = inputs.element_size()
+
+        elif isinstance(inputs, np.ndarray):
+            inputs_ = torch.from_numpy(inputs)
+            size, elem_bytes = list(inputs_.size()), inputs_.element_size()
+
+        elif isinstance(inputs, (list, tuple)):
+            size, elem_bytes = nested_list_size(inputs)
+            if batch_dim is not None and batch_dim < len(size):
+                size[batch_dim] = 1
+
+        else:
+            raise TypeError(
+                "Model contains a layer with an unsupported input or output type: "
+                f"{inputs}, type: {type(inputs)}"
+            )
+
+        return size, elem_bytes
+
+    @staticmethod
+    def get_param_count(
+        module: nn.Module, name: str, param: torch.Tensor
+    ) -> tuple[int, str]:
+        """
+        Get count of number of params, accounting for mask.
+
+        Masked models save parameters with the suffix "_orig" added.
+        They have a buffer ending with "_mask" which has only 0s and 1s.
+        If a mask exists, the sum of 1s in mask is number of params.
+        """
+        if name.endswith("_orig"):
+            without_suffix = name[:-5]
+            pruned_weights = rgetattr(module, f"{without_suffix}_mask")
+            if pruned_weights is not None:
+                parameter_count = int(torch.sum(pruned_weights))
+                return parameter_count, without_suffix
+        return param.nelement(), name
+
+    @staticmethod
+    def get_kernel_size(module: nn.Module) -> Union[int , list[int] , None]:
+        if hasattr(module, "kernel_size"):
+            k = module.kernel_size
+            kernel_size: Union[int , list[int]]
+            if isinstance(k, Iterable):
+                kernel_size = list(k)
+            elif isinstance(k, int):
+                kernel_size = int(k)
+            else:
+                raise TypeError(f"kernel_size has an unexpected type: {type(k)}")
+            return kernel_size
+        return None
+
+    def get_layer_name(self, show_var_name: bool, show_depth: bool) -> str:
+        layer_name = self.class_name
+        if show_var_name and self.var_name:
+            layer_name += f" ({self.var_name})"
+        if show_depth and self.depth > 0:
+            layer_name += f": {self.depth}"
+            if self.depth_index is not None:
+                layer_name += f"-{self.depth_index}"
+        return layer_name
+
+    def calculate_num_params(self) -> None:
+        """
+        Set num_params, trainable, inner_layers, and kernel_size
+        using the module's parameters.
+        """
+        self.num_params = 0
+        self.param_bytes = 0
+        self.trainable_params = 0
+        self.inner_layers = {}
+
+        final_name = ""
+        for name, param in self.module.named_parameters():
+            if is_lazy(param):
+                self.contains_lazy_param = True
+                continue
+            cur_params, name = self.get_param_count(self.module, name, param)
+
+            self.param_bytes += param.element_size() * cur_params
+
+            self.num_params += cur_params
+            if param.requires_grad:
+                self.trainable_params += cur_params
+
+            # kernel_size for inner layer parameters
+            ksize = list(param.size())
+            if name == "weight":
+                # to make [in_shape, out_shape, ksize, ksize]
+                if len(ksize) > 1:
+                    ksize[0], ksize[1] = ksize[1], ksize[0]
+
+            # RNN modules have inner weights such as weight_ih_l0
+            # Don't show parameters for the overall model, show for individual layers
+            if self.parent_info is not None or "." not in name:
+                self.inner_layers[name] = {
+                    ColumnSettings.KERNEL_SIZE: str(ksize),
+                    ColumnSettings.NUM_PARAMS: f"├─{cur_params:,}",
+                }
+                final_name = name
+        # Fix the final row to display more nicely
+        if self.inner_layers:
+            self.inner_layers[final_name][
+                ColumnSettings.NUM_PARAMS
+            ] = f"└─{self.inner_layers[final_name][ColumnSettings.NUM_PARAMS][2:]}"
+
+    def calculate_macs(self) -> None:
+        """
+        Set MACs using the module's parameters and layer's output size, which is
+        used for computing number of operations for Conv layers.
+
+        Please note: Returned MACs is the number of MACs for the full tensor,
+        i.e., taking the batch-dimension into account.
+        """
+        for name, param in self.module.named_parameters():
+            cur_params, name = self.get_param_count(self.module, name, param)
+            if name in ("weight", "bias"):
+                # ignore C when calculating Mult-Adds in ConvNd
+                if "Conv" in self.class_name:
+                    self.macs += int(
+                        cur_params * prod(self.output_size[:1] + self.output_size[2:])
+                    )
+                else:
+                    self.macs += self.output_size[0] * cur_params
+            # RNN modules have inner weights such as weight_ih_l0
+            elif "weight" in name or "bias" in name:
+                self.macs += prod(self.output_size[:2]) * cur_params
+
+    def check_recursive(self, layer_ids: set[int]) -> None:
+        """
+        If the current module is already-used, mark as (recursive).
+        Must check before adding line to the summary.
+        """
+        if self.layer_id in layer_ids:
+            self.is_recursive = True
+
+    def macs_to_str(self, reached_max_depth: bool) -> str:
+        """Convert MACs to string."""
+        if self.macs <= 0:
+            return "--"
+        if self.is_leaf_layer:
+            return f"{self.macs:,}"
+        if reached_max_depth:
+            sum_child_macs = sum(
+                child.macs for child in self.children if child.is_leaf_layer
+            )
+            return f"{sum_child_macs:,}"
+        return "--"
+
+    def num_params_to_str(self, reached_max_depth: bool) -> str:
+        """Convert num_params to string."""
+        if self.num_params == 0:
+            return "--"
+        if self.is_recursive:
+            return "(recursive)"
+        if reached_max_depth or self.is_leaf_layer:
+            param_count_str = f"{self.num_params:,}"
+            return param_count_str if self.trainable_params else f"({param_count_str})"
+        leftover_params = self.leftover_params()
+        return f"{leftover_params:,}" if leftover_params > 0 else "--"
+
+    def params_percent(
+        self, total_params: int, reached_max_depth: bool, precision: int = 2
+    ) -> str:
+        """Convert num_params to string."""
+        spacing = 5
+        zero = f"{' ' * spacing}--"
+        if total_params == 0:
+            return zero
+        if self.is_recursive:
+            return "(recursive)"
+        params = (
+            self.num_params
+            if reached_max_depth or self.is_leaf_layer
+            else self.leftover_params()
+        )
+        if params == 0:
+            return zero
+        return f"{params / total_params:>{precision + spacing}.{precision}%}"
+
+    def leftover_params(self) -> int:
+        """
+        Leftover params are the number of params this current layer has that are not
+        included in the child num_param counts.
+        """
+        return self.num_params - sum(
+            child.num_params if child.is_leaf_layer else child.leftover_params()
+            for child in self.children
+            if not child.is_recursive
+        )
+
+    def leftover_trainable_params(self) -> int:
+        return self.trainable_params - sum(
+            child.trainable_params
+            if child.is_leaf_layer
+            else child.leftover_trainable_params()
+            for child in self.children
+            if not child.is_recursive
+        )
+
+    def calculate_synapse_neuron_count(self, inputs: torch.Tensor, outputs: torch.Tensor) -> None:
+        # Only calculate the neuron and synapse counts for leafs (few basic torch layers).
+        # Recursive layers (nn.Modules consisting of basic layers will be updated later, once
+        # children are set up )
+        count_func = LayerParameterEventCalculator.get_calculate_synapse_neuron_count_for(
+            self.module.__class__.__name__)
+
+        # TODO: handle properly unrecognized layers
+        if count_func is None:
+            logging.warning(f"Couldn't calculate the synapses and neurons for layer with classname : "
+                            f"{self.module.__class__.__name__}")
+            self.synapse_count, self.neuron_count = None, None
+        else:
+            self.synapse_count, self.neuron_count = count_func(self, inputs, outputs)
+
+            # take into account the batch size : Each layer will have some number of trainable parameters, however,
+            # in hardware if batch size is not 1, they will have to be duplicated `batch_size` times
+            batch_size = inputs[0].shape[1]
+            self.synapse_count *= batch_size
+            self.neuron_count  *= batch_size
+
+    def calculate_events(self, inputs: torch.Tensor, outputs: torch.Tensor) -> None:
+
+        event_counter = LayerParameterEventCalculator.get_event_counter_for(self.module.__class__.__name__)
+        if event_counter is not None:
+            if self.spiking_events is None:
+                self.spiking_events = 0
+
+            if self.total_events is None:
+                self.total_events = 0
+
+            # get the number of spiking and total events for current timestep, for this layer
+            current_spiking_events, current_total_events = event_counter(self, inputs, outputs)
+
+            # update the spiking and total events
+            self.spiking_events += current_spiking_events
+            self.total_events += current_total_events
+
+            # record the number of passes through a particular layer
+            self.timestep += 1
+
+            # update the firing rate to mean across the timesteps firing rate
+            # here the firing rate is defined as spiking_events / total_events
+            self.firing_rate = self.spiking_events / self.total_events
+
+            # check if there is overriden energy contribution for this layer
+            custom_energy_contribution = LayerParameterEventCalculator.get_custom_energy_contribution_for(
+                self.module.__class__.__name__)
+
+            for i, deviceProfile in enumerate(self.device_profiles):
+
+                if custom_energy_contribution is not None:
+                    # execute the function for a Layer, taking particular 'inputs' and producing 'outputs', given
+                    # the device characteristic
+                    dEnergy = custom_energy_contribution(self, inputs, outputs, deviceProfile)
+                    self.total_energy_contributions[i] += dEnergy
+
+                # default loop : add to total energy per layer, the contribution from current timestep with
+                # recognizing whether layer is synaptic or not (which then multiplies energy of synapse operation or
+                # energy of neuron operation) and is the architecture truly spiking or not (if spiking architecture,
+                # add only the spiking events, if not, add energy contribution, even if value produced is not a spike)
+                else:
+                    # if layer is synaptic, get the energy per synapse event to add to total contribution, otherwise
+                    # get the energy per neuron event
+                    energy_per_event = deviceProfile.energy_per_neuron_event
+                    if self.is_synaptic:
+                        energy_per_event = deviceProfile.energy_per_synapse_event
+
+                    # if spiking architecture, add only energy when spike was generated
+                    if deviceProfile.is_spiking_architecture:
+                        self.total_energy_contributions[i] += current_spiking_events * energy_per_event
+
+                    # not a spiking architecture, always add energy
+                    else:
+                        self.total_energy_contributions[i] += current_total_events * energy_per_event
+
+
+def nested_list_size(inputs: Union[Sequence[Any] , torch.Tensor]) -> tuple[list[int], int]:
+    """Flattens nested list size."""
+    if hasattr(inputs, "tensors"):
+        size, elem_bytes = nested_list_size(inputs.tensors)
+    elif isinstance(inputs, torch.Tensor):
+        size, elem_bytes = list(inputs.size()), inputs.element_size()
+    elif isinstance(inputs, np.ndarray):
+        inputs_torch = torch.from_numpy(inputs)  # preserves dtype
+        size, elem_bytes = list(inputs_torch.size()), inputs_torch.element_size()
+    elif not hasattr(inputs, "__getitem__") or not inputs:
+        size, elem_bytes = [], 0
+    elif isinstance(inputs, dict):
+        size, elem_bytes = nested_list_size(list(inputs.values()))
+    elif (
+        hasattr(inputs, "size")
+        and callable(inputs.size)
+        and hasattr(inputs, "element_size")
+        and callable(inputs.element_size)
+    ):
+        size, elem_bytes = list(inputs.size()), inputs.element_size()
+    elif isinstance(inputs, (list, tuple)):
+        size, elem_bytes = nested_list_size(inputs[0])
+    else:
+        size, elem_bytes = [], 0
+
+    return size, elem_bytes
+
+
+def prod(num_list: Union[Iterable[int] , torch.Size]) -> int:
+    result = 1
+    if isinstance(num_list, Iterable):
+        for item in num_list:
+            result *= prod(item) if isinstance(item, Iterable) else item
+    return result
+
+
+def rgetattr(module: nn.Module, attr: str) -> Union[torch.Tensor , None]:
+    """Get the tensor submodule called attr from module."""
+    for attr_i in attr.split("."):
+        if not hasattr(module, attr_i):
+            return None
+        module = getattr(module, attr_i)
+    assert isinstance(module, torch.Tensor)
+    return module
+
+
+def get_children_layers(summary_list: list[LayerInfo], index: int) -> list[LayerInfo]:
+    """Fetches all of the children of a given layer."""
+    num_children = 0
+    for layer in summary_list[index + 1:]:
+        if layer.depth <= summary_list[index].depth:
+            break
+        num_children += 1
+    return summary_list[index + 1: index + 1 + num_children]
diff --git a/snntorch/energy_estimation/layer_parameter_event_calculator.py b/snntorch/energy_estimation/layer_parameter_event_calculator.py
new file mode 100644
index 00000000..467e9112
--- /dev/null
+++ b/snntorch/energy_estimation/layer_parameter_event_calculator.py
@@ -0,0 +1,197 @@
+import torch
+import torch.nn as nn
+from typing import Callable, Tuple, Dict, TypeVar, ClassVar, List, Type, Union
+from .device_profile import DeviceProfile
+from copy import deepcopy
+
+import snntorch
+from .utils import *
+
+# signature for functions
+synapse_neuron_count_signature = TypeVar("synapse_neuron_count_signature",
+                                         bound=Callable[[object, torch.Tensor, torch.Tensor], Tuple[int, int]])
+#                                                       TODO: passing object is not a good solution probably in typing
+#                                                        This most likely should be LayerInfo (or some base class of it)
+event_counter_signature = TypeVar("event_counter_signature",
+                                  bound=Callable[[object, torch.Tensor, torch.Tensor], Tuple[int, int]])
+#                                                       TODO: passing object is not a good solution probably in typing
+#                                                        This most likely should be LayerInfo (or some base class of it)
+
+custom_energy_contribution_signature = TypeVar("custom_energy_contribution_signature",
+                                               bound=Callable[[object, torch.Tensor, torch.Tensor, DeviceProfile],
+                                               float])
+#                                                       TODO: passing object is not a good solution probably in typing
+#                                                        This most likely should be LayerInfo (or some base class of it)
+
+
+# calculate the neuron and synapse count for Linear layer
+def synapse_neuron_count_for_linear(layer_info, inputs: torch.Tensor, outputs: torch.Tensor) -> Tuple[int, int]:
+    if len(inputs) == 1:
+        val = torch.Tensor(inputs[0])
+    else:
+        val = torch.Tensor(inputs)
+
+    return layer_info.trainable_params, 0
+
+
+# calculate the number of events (spiking/total) for a Linear layer
+def count_events_for_linear(layer_info, inputs: torch.Tensor, outputs: torch.Tensor) -> Tuple[int, int]:
+    current_spiking_events = (float(torch.sum(inputs[0]).detach().cpu()) + layer_info.include_bias_in_events) * \
+                             prod(outputs[0].shape)
+    current_total_events = (prod(inputs[0].size()) + layer_info.include_bias_in_events) * prod(outputs[0].shape)
+
+    # in the above calculations the is included twice (these values scales quadratically rather than linearly)
+    # to solve this, get the batch size value (shape[1] or shape[0], depending on whether
+    # network_requires_first_dim_as_time is True or False ) and divide both values by it
+
+    if layer_info.network_requires_first_dim_as_time:
+        batch_size = inputs[0].shape[1]
+    else:
+        batch_size = inputs[0].shape[0]
+
+    return int(current_spiking_events / batch_size), int(current_total_events / batch_size)
+
+
+# calculate the neuron and synapse count for snntorch activation function ( like Leaky)
+def synapse_neuron_count_for_snn_neuron(layer_info, inputs: torch.Tensor, outputs: torch.Tensor) -> Tuple[int, int]:
+    val = torch.Tensor(outputs[0])
+    layer_info.is_synaptic = False
+
+    if layer_info.network_requires_first_dim_as_time:
+        return 0, prod(outputs[0].shape[2:])
+    return 0, prod(outputs[0].shape[1:])
+
+
+# calculate the number of events (spiking/total) for a snntorch activation function ( like Leaky)
+def count_events_for_snn_neuron(layer_info, inputs: torch.Tensor, outputs: torch.Tensor) -> Tuple[int, int]:
+    current_spiking_events = int(torch.sum(outputs[0].detach()).cpu())
+    current_total_events = prod(outputs[0].size())
+    return current_spiking_events, current_total_events
+
+
+# calculate the neuron and synapse count for the convolution layers
+def synapse_neuron_count_for_conv(layer_info, inputs: torch.Tensor, outputs: torch.Tensor) -> Tuple[int, int]:
+    # TODO : make sure that's correct
+    if len(inputs) == 1:
+        val = torch.Tensor(inputs[0])
+    else:
+        val = torch.Tensor(inputs)
+    return layer_info.trainable_params, 0
+
+
+# calculate the number of events (spiking/total) for a convolution layers
+def count_events_for_conv(layer_info, inputs: torch.Tensor, outputs: torch.Tensor) -> Tuple[int, int]:
+    connections = layer_info.trainable_params - (1 - layer_info.include_bias_in_events) * prod(
+        layer_info.module.bias.shape)
+
+    # get input and output channels
+    in_channels, out_channels, kernel_size, bias_size = (layer_info.module.in_channels, layer_info.module.out_channels,
+                                                         layer_info.module.kernel_size,
+                                                         layer_info.module.bias.data.shape)
+
+    current_total_events = (in_channels * prod(outputs.shape) *
+                            (prod(kernel_size) + layer_info.include_bias_in_events * prod(bias_size)))
+
+    # TODO : seems like a very ugly solution, maybe make it more elegant ?
+    ones_conv = deepcopy(layer_info.module)
+    # remove all the hooks from a copy of this layer
+    ones_conv._forward_hooks = {}
+    ones_conv.weight.data = torch.ones_like(ones_conv.weight.data)
+
+    if layer_info.include_bias_in_events:
+        ones_conv.bias.data = torch.ones_like(ones_conv.bias.data)
+    else:
+        ones_conv.bias.data = torch.zeros_like(ones_conv.bias.data)
+
+    with torch.no_grad():
+        current_spiking_events = int(torch.sum(ones_conv(inputs[0]).cpu()))
+
+    return current_spiking_events, current_total_events
+
+
+class LayerParameterEventCalculator(object):
+    """
+        Singleton which provides an API for getting neuron/synapse/event counts for the estimate_energy.
+        Every basic pytorch layer (like Linear, Conv2d or activation function ) has a small function which will return
+        to synapse/neuron/event counts, which will be calculated in estimate_energy.
+
+        Every layer can also have a `custom_energy_contribution` function (with signature
+        `custom_energy_contribution_signature`), which will override default calculations (multiply the number of
+        synapses and neurons by energies required to perform these operations ), and developer can specify custom
+        energy value (based on input and output tensors + significant amount of information from LayerInfo)
+
+        An example of adding a layer using this singleton :
+
+        >>> class ANewLinearLayer(nn.module):
+        >>>     def forward(self, x : torch.Tensor):
+        >>>        # ...
+        >>>
+        >>>
+        >>>  LayerParameterEventCalculator.register_new_layer(ANewLinearLayer,
+        >>>                                             synapse_neuron_count_for_linear,
+        >>>                                             count_events_for_linear)
+
+    """
+
+    # a dictionary mapping a layer class name (like Linear) to tuple, where :
+    # - first element is a function returning 2 ints representing number of synapses and neurons (function should
+    # have signature `synapse_neuron_count_signature`)
+    # - second element is function returning 2 ints representing number of spiking and total events (function should
+    # have signature 'event_counter_signature')
+    # - third optional element returning a float representing energy required to perform the operation (function, if
+    # specified should have signature custom_energy_contribution_signature )
+    _recognized_layers: Dict[str, Tuple[
+        synapse_neuron_count_signature, event_counter_signature, Union[custom_energy_contribution_signature, None]]] = \
+        {
+            nn.Linear.__name__: (synapse_neuron_count_for_linear, count_events_for_linear, None),
+            snntorch.Leaky.__name__: (synapse_neuron_count_for_snn_neuron, count_events_for_snn_neuron, None),
+            nn.Conv2d.__name__: (synapse_neuron_count_for_conv, count_events_for_conv, None),
+            nn.Conv1d.__name__: (synapse_neuron_count_for_conv, count_events_for_conv, None)
+        }
+
+    @staticmethod
+    def register_new_layer(layer: Type[torch.nn.Module],
+                           synapse_neuron_count_lambda: synapse_neuron_count_signature,
+                           event_counter_lambda: event_counter_signature,
+                           custom_energy_contribution: custom_energy_contribution_signature = None,
+                           override_layer_info: bool = False):
+
+        if layer.__class__.__name__ in LayerParameterEventCalculator._recognized_layers:
+            if override_layer_info:
+                pass
+            else:
+                raise Exception("overriding layer information, and override_layer_info was set to False!")
+
+        LayerParameterEventCalculator._recognized_layers[layer.__name__] = (synapse_neuron_count_lambda,
+                                                                            event_counter_lambda,
+                                                                            custom_energy_contribution)
+
+    @staticmethod
+    def get_calculate_synapse_neuron_count_for(name: str,
+                                               ignore_not_recognized: bool = True) -> Union[synapse_neuron_count_signature , None]:
+
+        if name not in LayerParameterEventCalculator._recognized_layers:
+            if ignore_not_recognized:
+                return None
+            else:
+                raise Exception(f"couldn't get the function to calculate the synapse neuron count for layer with "
+                                f"classname :\'{name}\'")
+
+        return LayerParameterEventCalculator._recognized_layers[name][0]
+
+    @staticmethod
+    def get_event_counter_for(name: str, ignore_not_recognized: bool = True) -> Union[event_counter_signature , None]:
+        if name not in LayerParameterEventCalculator._recognized_layers:
+            if ignore_not_recognized:
+                return None
+            else:
+                raise Exception(f"couldn't get the function to count the total and spiking events for layer with "
+                                f"classname :\'{name}\'")
+
+        return LayerParameterEventCalculator._recognized_layers[name][1]
+
+    @staticmethod
+    def get_custom_energy_contribution_for(name: str) -> Union[custom_energy_contribution_signature , None]:
+        if name not in LayerParameterEventCalculator._recognized_layers:
+            return None
+        return LayerParameterEventCalculator._recognized_layers[name][2]
diff --git a/snntorch/energy_estimation/model_statistics.py b/snntorch/energy_estimation/model_statistics.py
new file mode 100644
index 00000000..0b9e945e
--- /dev/null
+++ b/snntorch/energy_estimation/model_statistics.py
@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from typing import Any, List
+from .enums import Units
+from .FormattingOptions import CONVERSION_FACTORS, FormattingOptions
+from .layer_info import LayerInfo
+from .device_profile import DeviceProfile
+
+"""
+    Model statistics has been taken from torchinfo project https://github.com/TylerYep/torchinfo
+"""
+
+class ModelStatistics:
+    """Class for storing results of the summary."""
+
+    def __init__(
+        self,
+        summary_list: list[LayerInfo],
+        input_size: Any,
+        total_input_size: int,
+        formatting: FormattingOptions,
+        device_profiles : List[DeviceProfile]
+    ) -> None:
+        self.summary_list = summary_list
+        self.input_size = input_size
+        self.formatting = formatting
+        self.total_input = total_input_size
+        self.total_mult_adds = 0
+        self.total_params, self.trainable_params = 0, 0
+        self.total_param_bytes, self.total_output_bytes = 0, 0
+        self.total_energy = 0
+        self.device_profiles = device_profiles
+        self.total_energies = self.summary_list[0].total_energy_contributions
+
+        # TODO: Figure out why the below functions using max() are ever 0
+        # (they should always be non-negative), and remove the call to max().
+        # Investigation: https://github.com/TylerYep/torchinfo/pull/195
+        for layer_info in summary_list:
+            if layer_info.is_leaf_layer:
+                self.total_mult_adds += layer_info.macs
+                if layer_info.num_params > 0:
+                    # x2 for gradients
+                    self.total_output_bytes += layer_info.output_bytes * 2
+                if layer_info.is_recursive:
+                    continue
+                self.total_params += max(layer_info.num_params, 0)
+                self.total_param_bytes += layer_info.param_bytes
+                self.trainable_params += max(layer_info.trainable_params, 0)
+            else:
+                if layer_info.is_recursive:
+                    continue
+                leftover_params = layer_info.leftover_params()
+                leftover_trainable_params = layer_info.leftover_trainable_params()
+                self.total_params += max(leftover_params, 0)
+                self.trainable_params += max(leftover_trainable_params, 0)
+        self.formatting.set_layer_name_width(summary_list)
+
+    def __repr__(self) -> str:
+        """Print results of the summary."""
+        divider = "=" * self.formatting.get_total_width()
+        total_params = ModelStatistics.format_output_num(
+            self.total_params, self.formatting.params_units
+        )
+        trainable_params = ModelStatistics.format_output_num(
+            self.trainable_params, self.formatting.params_units
+        )
+        non_trainable_params = ModelStatistics.format_output_num(
+            self.total_params - self.trainable_params, self.formatting.params_units
+        )
+
+        all_layers = self.formatting.layers_to_str(self.summary_list, self.total_params)
+
+        total_energy_strings = []
+        for device_profile, total_energy in zip(self.device_profiles, self.total_energies):
+            total_energy_strings.append(f"Total energy for device [{str(device_profile)}] : {total_energy} J/inf\n")
+
+        summary_str = (
+            f"{divider}\n"
+            f"{self.formatting.header_row()}{divider}\n"
+            f"{all_layers}{divider}\n"
+        )
+        summary_str += "".join(total_energy_strings)
+
+        return summary_str
+
+    @staticmethod
+    def float_to_megabytes(num: int) -> float:
+        """Converts a number (assume floats, 4 bytes each) to megabytes."""
+        return num * 4 / 1e6
+
+    @staticmethod
+    def to_megabytes(num: int) -> float:
+        """Converts bytes to megabytes."""
+        return num / 1e6
+
+    @staticmethod
+    def to_readable(num: int, units: Units = Units.AUTO) -> tuple[Units, float]:
+        """Converts a number to millions, billions, or trillions."""
+        if units == Units.AUTO:
+            if num >= 1e12:
+                return Units.TERABYTES, num / 1e12
+            if num >= 1e9:
+                return Units.GIGABYTES, num / 1e9
+            return Units.MEGABYTES, num / 1e6
+        return units, num / CONVERSION_FACTORS[units]
+
+    @staticmethod
+    def format_output_num(num: int, units: Units) -> str:
+        units_used, converted_num = ModelStatistics.to_readable(num, units)
+        if converted_num.is_integer():
+            converted_num = int(converted_num)
+        units_display = "" if units_used == Units.NONE else f" ({units_used})"
+        fmt = "d" if isinstance(converted_num, int) else ".2f"
+        return f"{units_display}: {converted_num:,{fmt}}"
diff --git a/snntorch/energy_estimation/utils.py b/snntorch/energy_estimation/utils.py
new file mode 100644
index 00000000..c3521b89
--- /dev/null
+++ b/snntorch/energy_estimation/utils.py
@@ -0,0 +1,10 @@
+from typing import Iterable, Union
+import torch
+
+
+def prod(num_list: Union[Iterable[int] , torch.Size]) -> int:
+    result = 1
+    if isinstance(num_list, Iterable):
+        for item in num_list:
+            result *= prod(item) if isinstance(item, Iterable) else item
+    return result


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index eddf1d4d..6095dd5f 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -2,7 +2,9 @@
 import snntorch as snn
 import torch
 import torch.nn as nn
-
+import math
+from snntorch.energy_estimation.energy_estimation_network_interface import EnergyEstimationNetworkInterface
+from typing import Union
 
 # class Net(nn.Module):
 #     def __init__(self):
@@ -38,4 +40,368 @@ def forward(self, x):
         return self.lif1.spk, self.lif1.mem
 
 
-# net = Net().to(device)
+class EnergyEfficiencyNetTest1(nn.Module):
+    def __init__(self, beta: float, num_timesteps: Union[int , None] = None):
+        super().__init__()
+        self.fc1 = nn.Linear(1, 1)
+        self.lif1 = snn.Leaky(beta=beta)
+        self.fc2 = nn.Linear(1, 1)
+        self.lif2 = snn.Leaky(beta=beta)
+
+        self.num_timesteps = num_timesteps
+
+    def forward(self, x: torch.Tensor, num_timesteps: Union[int, None] = None):
+        assert num_timesteps is not None or self.num_timesteps is not None, "please specify the number of timestep for this testing network!"
+
+        # specifying timesteps in forward will override the self.num_timesteps
+        if num_timesteps is not None:
+            timesteps = num_timesteps
+        else:
+            timesteps = self.num_timesteps
+
+        assert x.size(0) == timesteps
+
+        # Initialize the hidden states of LIFs
+        mem1 = self.lif1.init_leaky()
+        mem2 = self.lif2.init_leaky()
+
+        # final layers values in time
+        spk2_rec, mem2_rec = [], []
+
+        for timestep in range(timesteps):
+            cur1 = self.fc1(x[timestep])
+            spk1, mem1 = self.lif1(cur1, mem1)
+            cur2 = self.fc2(spk1)
+            spk2, mem2 = self.lif2(cur2, mem2)
+
+            spk2_rec.append(spk2)
+            mem2_rec.append(mem2)
+
+        return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)
+
+
+class EnergyEfficiencyNetTest2(nn.Module):
+    def __init__(self, beta: float):
+        super().__init__()
+        self.fc1 = nn.Linear(32, 64)
+        self.lif1 = snn.Leaky(beta=beta, threshold=0.2)
+        self.fc2 = nn.Linear(64, 10)
+        self.lif2 = snn.Leaky(beta=beta, threshold=0.4)
+        self.reset()
+
+    def forward(self, x: torch.Tensor):
+        cur1 = self.fc1(x)
+        spk1, self.mem1 = self.lif1(cur1, self.mem1)
+        cur2 = self.fc2(spk1)
+        spk2, self.mem2 = self.lif2(cur2, self.mem2)
+
+        return spk2, self.mem2
+
+    def reset(self):
+        # Initialize the hidden states of LIFs
+        self.mem1 = self.lif1.init_leaky()
+        self.mem2 = self.lif2.init_leaky()
+
+
+class EnergyEfficiencyNetTest3(EnergyEstimationNetworkInterface):
+    def __init__(self, beta: float):
+        super().__init__()
+        self.fc1 = nn.Linear(1, 1)
+        self.lif1 = snn.Leaky(beta=beta, threshold=0.5)
+        self.fc2 = nn.Linear(1, 1)
+        self.lif2 = snn.Leaky(beta=beta, threshold=0.25)
+        self.reset()
+
+    def forward(self, x: torch.Tensor):
+        cur1 = self.fc1(x)
+        spk1, self.mem1 = self.lif1(cur1, self.mem1)
+        cur2 = self.fc2(spk1)
+        spk2, self.mem2 = self.lif2(cur2, self.mem2)
+
+        return spk2, self.mem2
+
+    def forward_full(self, x: torch.Tensor):
+        cur1 = self.fc1(x)
+        spk1, self.mem1 = self.lif1(cur1, self.mem1)
+        cur2 = self.fc2(spk1)
+        spk2, self.mem2 = self.lif2(cur2, self.mem2)
+
+        return spk1, self.mem1, spk2, self.mem2
+
+    def reset(self):
+        # Initialize the hidden states of LIFs
+        self.mem1 = self.lif1.init_leaky()
+        self.mem2 = self.lif2.init_leaky()
+
+
+class EnergyEfficiencyNetTest4(EnergyEstimationNetworkInterface):
+    def __init__(self):
+        super().__init__()
+        self.cnn1 = nn.Conv2d(1, 1, kernel_size=3, padding=0)
+        self.lif1 = snn.Leaky(beta=0.5, threshold=1)
+        self.fc2 = nn.Linear(9, 1)
+        self.lif2 = snn.Leaky(beta=0.25, threshold=1)
+        self.reset()
+
+    def forward(self, x: torch.Tensor):
+        cur1 = self.cnn1(x)
+        spk1, self.mem1 = self.lif1(cur1, self.mem1)
+
+        # reshape into (time, batchsize)
+        spk1 = spk1.reshape(spk1.shape[0], -1)
+        cur2 = self.fc2(spk1)
+        spk2, self.mem2 = self.lif2(cur2, self.mem2)
+
+        return spk2, self.mem2
+
+    def forward_full(self, x: torch.Tensor):
+        cur1 = self.cnn1(x)
+        spk1, self.mem1 = self.lif1(cur1, self.mem1)
+        spk1 = spk1.reshape(spk1.shape[0], -1)
+        cur2 = self.fc2(spk1)
+        spk2, self.mem2 = self.lif2(cur2, self.mem2)
+        return spk1, self.mem1, spk2, self.mem2
+
+    def reset(self):
+        # Initialize the hidden states of LIFs
+        self.mem1 = self.lif1.init_leaky()
+        self.mem2 = self.lif2.init_leaky()
+
+
+class EnergyEfficiencyNetTest5Large(EnergyEstimationNetworkInterface):
+    def __init__(self):
+        super().__init__()
+        # added just to remove a warning about declaration of variables outsides __init__
+        self.mem1 = self.mem2 = self.mem3 = self.mem4 = self.mem5 = None
+
+        self.cnn1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
+        self.lif1 = snn.Leaky(beta=0.5, threshold=1)
+
+        self.cnn2 = nn.Conv2d(16, 16, kernel_size=5, dilation=2)
+        self.lif2 = snn.Leaky(beta=0.5, threshold=1)
+
+        self.cnn3 = nn.Conv2d(16, 16, kernel_size=7, dilation=3)
+        self.lif3 = snn.Leaky(beta=0.5, threshold=1)
+
+        self.fc1 = nn.Linear(576, 32)
+        self.lif4 = snn.Leaky(beta=0.5, threshold=1)
+
+        self.fc2 = nn.Linear(32, 10)
+        self.lif5 = snn.Leaky(beta=0.5, threshold=1)
+        self.reset()
+
+    def forward(self, x: torch.Tensor):
+        cur1 = self.cnn1(x)
+        spk1, self.mem1 = self.lif1(cur1, self.mem1)
+
+        cur2 = self.cnn2(spk1)
+        spk2, self.mem2 = self.lif2(cur2, self.mem2)
+
+        cur3 = self.cnn3(spk2)
+        spk3, self.mem3 = self.lif3(cur3, self.mem3)
+
+        # flatten it out
+        spk3 = spk3.reshape(spk1.shape[0], 16 * 6 * 6)
+
+        cur4 = self.fc1(spk3)
+        spk4, self.mem4 = self.lif4(cur4, self.mem4)
+
+        cur5 = self.fc2(spk4)
+        spk5, self.mem5 = self.lif5(cur5, self.mem5)
+
+        return spk5, self.mem5
+
+    def reset(self):
+        # Initialize the hidden states of LIFs
+        self.mem1 = self.lif1.init_leaky()
+        self.mem2 = self.lif2.init_leaky()
+        self.mem3 = self.lif3.init_leaky()
+        self.mem4 = self.lif4.init_leaky()
+        self.mem5 = self.lif5.init_leaky()
+
+
+class EnergyEfficiencyNet1DConv(EnergyEstimationNetworkInterface):
+    def __init__(self):
+        super().__init__()
+        # added just to remove a warning about declaration of variables outsides __init__
+        self.mem1 = self.mem2 = self.mem3 = self.mem4 = None
+        self.cnn1 = nn.Conv1d(1, 16, kernel_size=3, padding=1)
+        self.lif1 = snn.Leaky(beta=0.5, threshold=1)
+
+        self.cnn2 = nn.Conv1d(16, 16, kernel_size=5)
+        self.lif2 = snn.Leaky(beta=0.5, threshold=1)
+
+        self.reshape_val = 16 * 28
+        self.fc1 = nn.Linear(self.reshape_val, 64)
+        self.lif3 = snn.Leaky(beta=0.5, threshold=1)
+
+        self.cls = nn.Linear(64, 1)
+        self.cls_activ = snn.Leaky(beta=0.5, threshold=1)
+        self.reset()
+
+    def forward(self, x: torch.Tensor):
+        cur1 = self.cnn1(x)
+        spk1, self.mem1 = self.lif1(cur1, self.mem1)
+
+        cur2 = self.cnn2(spk1)
+        spk2, self.mem2 = self.lif2(cur2, self.mem2)
+        spk2 = spk2.reshape(-1, self.reshape_val)
+
+        cur3 = self.fc1(spk2)
+        spk3, self.mem3 = self.lif3(cur3, self.mem3)
+
+        cur4 = self.cls(spk3)
+        spk4, self.mem4 = self.cls_activ(cur4, self.mem4)
+        return spk4, self.mem4
+
+    def reset(self):
+        # Initialize the hidden states of LIFs
+        self.mem1 = self.lif1.init_leaky()
+        self.mem2 = self.lif2.init_leaky()
+        self.mem3 = self.lif3.init_leaky()
+        self.mem4 = self.cls_activ.init_leaky()
+
+
+class UnrecognizedLeafTestLayer(nn.Module):
+    __constants__ = ['in_features', 'out_features']
+    in_features: int
+    out_features: int
+    weight: torch.Tensor
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None) -> None:
+        factory_kwargs = {'device': device, 'dtype': dtype}
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs))
+        if bias:
+            self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return nn.functional.linear(input, self.weight, self.bias)
+
+
+class EnergyEfficiencyNetUnrecognizedLayers(EnergyEstimationNetworkInterface):
+    def __init__(self):
+        super().__init__()
+        self.mem1 = self.mem2 = None
+        self.fc1 = UnrecognizedLeafTestLayer(128, 32)
+        self.lif1 = snn.Leaky(beta=0.5, threshold=1)
+        self.fc2 = UnrecognizedLeafTestLayer(32, 16)
+        self.lif2 = snn.Leaky(beta=0.5, threshold=1)
+        self.reset()
+
+    def forward(self, x: torch.Tensor):
+        cur1 = self.fc1(x)
+        spk1, self.mem1 = self.lif1(cur1, self.mem1)
+        cur2 = self.fc2(spk1)
+        spk2, self.mem2 = self.lif2(cur2, self.mem2)
+
+        return spk2, self.mem2
+
+    def reset(self):
+        # Initialize the hidden states of LIFs
+        self.mem1 = self.lif1.init_leaky()
+        self.mem2 = self.lif2.init_leaky()
+
+
+class NestedNetworkModule(nn.Module):
+    def __init__(self, inp: int, out: int):
+        super().__init__()
+        self.fc1 = nn.Linear(inp, out)
+        self.lif1 = snn.Leaky(beta=1)
+        self.mem1 = None
+        self.reset()
+
+    def forward(self, x):
+        x = self.fc1(x)
+        spk1, self.mem1 = self.lif1(x, self.mem1)
+        return spk1
+
+    def reset(self):
+        self.mem1 = self.lif1.init_leaky()
+
+
+class NestedNetworkBlock(nn.Module):
+    def __init__(self, inp, hidden1, hidden2, out):
+        super().__init__()
+        self.module1 = NestedNetworkModule(inp, hidden1)
+        self.module2 = NestedNetworkModule(hidden1, hidden2)
+        self.module3 = NestedNetworkModule(hidden2, out)
+
+    def forward(self, x):
+        x = self.module1(x)
+        x = self.module2(x)
+        x = self.module3(x)
+        return x
+
+    def reset(self):
+        self.module1.reset()
+        self.module2.reset()
+        self.module3.reset()
+
+
+class EnergyEfficiencyNestedNetworkTest(EnergyEstimationNetworkInterface):
+    def __init__(self):
+        super().__init__()
+        self.encoder = NestedNetworkBlock(128, 64, 32, 16)
+        self.decoder = NestedNetworkBlock(16, 32, 64, 128)
+
+    def forward(self, x: torch.Tensor):
+        x = self.encoder(x)
+        return self.decoder(x)
+
+    def reset(self):
+        # Initialize the hidden states of LIFs
+        self.encoder.reset()
+        self.decoder.reset()
+
+
+class EnergyEfficiencyNestedNetworkTestUnfolded(EnergyEstimationNetworkInterface):
+
+    def __init__(self, Beta=0.9):
+        super().__init__()
+        self.fc1 = nn.Linear(128, 64)
+        self.lif1 = snn.Leaky(Beta)
+
+        self.fc2 = nn.Linear(64, 32)
+        self.lif2 = snn.Leaky(Beta)
+
+        self.fc3 = nn.Linear(32, 16)
+        self.lif3 = snn.Leaky(Beta)
+
+        self.fc4 = nn.Linear(16, 32)
+        self.lif4 = snn.Leaky(Beta)
+
+        self.fc5 = nn.Linear(32, 64)
+        self.lif5 = snn.Leaky(Beta)
+
+        self.fc6 = nn.Linear(64, 128)
+        self.lif6 = snn.Leaky(Beta)
+        self.reset()
+
+    def forward(self, x: torch.Tensor):
+        spk, self.mem1 = self.lif1(self.fc1(x), self.mem1)
+        spk, self.mem2 = self.lif2(self.fc2(spk), self.mem2)
+        spk, self.mem3 = self.lif3(self.fc3(spk), self.mem3)
+        spk, self.mem4 = self.lif4(self.fc4(spk), self.mem4)
+        spk, self.mem5 = self.lif5(self.fc5(spk), self.mem5)
+        return self.lif6(self.fc6(spk), self.mem6)
+
+    def reset(self):
+        self.mem1 = self.lif1.init_leaky()
+        self.mem2 = self.lif1.init_leaky()
+        self.mem3 = self.lif1.init_leaky()
+        self.mem4 = self.lif1.init_leaky()
+        self.mem5 = self.lif1.init_leaky()
+        self.mem6 = self.lif1.init_leaky()
diff --git a/tests/test_energy_estimation.py b/tests/test_energy_estimation.py
new file mode 100644
index 00000000..fbbe1580
--- /dev/null
+++ b/tests/test_energy_estimation.py
@@ -0,0 +1,1125 @@
+import numpy as np
+import pytest
+from snntorch.energy_estimation.estimate_energy import *
+from snntorch.energy_estimation.device_profile_registry import *
+from snntorch.energy_estimation.layer_parameter_event_calculator import (LayerParameterEventCalculator,
+                                                                         synapse_neuron_count_for_linear,
+                                                                         count_events_for_linear)
+from tests.conftest import *
+import torch
+
+
+def test_estimate_energy_network1():
+    """
+        Simplest network test :
+            - non-neuromorphic architecture (cpu) with energy 8.6e-9 per synapse/neuron event
+
+            - 2 Layers, simplest architecture : ( not including bias in calculations )
+
+                    self.fc1 = nn.Linear(1, 1)              # 1 synaptic event (2 with bias)
+                    self.lif1 = snn.Leaky(beta=beta)        # 1 neuron event
+                    self.fc2 = nn.Linear(1, 1)              # 1 synaptic event (2 with bias)
+                    self.lif2 = snn.Leaky(beta=beta)        # 1 neuron event
+
+        Tests :
+            - assert that energy returned is to expected value (without bias 4 * 8.6e-9 = 3.44e-8, with bias 5.16e-8)
+            - assert that the synapse/neuron/total counts are as expected
+    """
+
+    # Prepare the input data
+    # here
+    num_values = 1
+    percentage_of_ones = 0.
+    ones_num = round(percentage_of_ones * num_values)
+    zeros_num = num_values - ones_num
+    inp_test = torch.Tensor([[[1]]] * ones_num + [[[0]]] * zeros_num)
+
+    # create the network (the "device under test")
+    v = EnergyEfficiencyNetTest1(beta=0.9, num_timesteps=inp_test.size(0))
+
+    # register the device
+    energy_per_synapse_event = 8.6e-9
+    energy_per_neuron_event = 8.6e-9
+    DeviceProfileRegistry.add_device("cpu-test1", energy_per_synapse_event,
+                                     energy_per_neuron_event,
+                                     False)
+
+    # get the summary
+    summary_list = estimate_energy(model=v, input_data=inp_test, devices="cpu-test1",
+                                   include_bias_term_in_events=False)
+    # Expected energy values
+    expected_value_without_bias = 2 * energy_per_synapse_event + 2 * energy_per_neuron_event  # 3.44e-8
+    expected_value_with_bias = 4 * energy_per_synapse_event + 2 * energy_per_neuron_event  # 5.16e-8
+
+    # get the energies
+    energies = summary_list.total_energies
+
+    # only one device specified, therefore length of energies should be 1, and equal to `expected_value_without_bias`,
+    # assert that with 1% tolerance
+    assert len(energies) == 1
+    assert pytest.approx(expected_value_without_bias, rel=0.01) == energies[0]
+
+    # assert the expected numbers of synapses / neurons / total events
+    # topmost layer, the EnergyEfficiencyNetTest1
+    assert summary_list.summary_list[0].total_events == 4  # 4 total events
+    assert summary_list.summary_list[0].synapse_count == 4  # 4 total synapses (2 for weights, 2 for biases)
+    assert summary_list.summary_list[0].neuron_count == 2  # 2 neuron events
+
+    assert summary_list.summary_list[1].total_events == 1  # 1 synapse event (bias is not included)
+    assert summary_list.summary_list[1].synapse_count == 2  # 2 synapses (single weight + bias)
+    assert summary_list.summary_list[1].neuron_count == 0  # 0 neurons (layer is Linear (synaptic), no neurons)
+
+    assert summary_list.summary_list[2].total_events == 1  # 1 synapse event (bias is not included)
+    assert summary_list.summary_list[2].synapse_count == 0  # 0 synapses (layer is neuron, no synapses)
+    assert summary_list.summary_list[2].neuron_count == 1  # 1 neurons (layer is a neuron, output shape is 1, one event)
+
+    assert summary_list.summary_list[3].total_events == 1  # 1 synapse event (bias is not included)
+    assert summary_list.summary_list[3].synapse_count == 2  # 2 synapses (single weight + bias)
+    assert summary_list.summary_list[3].neuron_count == 0  # 0 neurons (layer is Linear (synaptic), no neurons)
+
+    assert summary_list.summary_list[4].total_events == 1  # 1 synapse event (bias is not included)
+    assert summary_list.summary_list[4].synapse_count == 0  # 0 synapses (layer is neuron, no synapses)
+    assert summary_list.summary_list[4].neuron_count == 1  # 1 neurons (layer is a neuron, output shape is 1, one event)
+
+    # now get the summary with bias included in it's calculations
+    summary_list_bias = estimate_energy(model=v, input_data=inp_test, devices="cpu-test1",
+                                        include_bias_term_in_events=True)
+    energies = summary_list_bias.total_energies
+
+    # assert now the same thing, but when bias is included
+    assert summary_list_bias.summary_list[0].total_events == 6  # 4 total events
+    assert summary_list_bias.summary_list[0].synapse_count == 4  # 4 total synapses (2 for weights, 2 for biases)
+    assert summary_list_bias.summary_list[0].neuron_count == 2  # 2 neuron events
+
+    assert summary_list_bias.summary_list[1].total_events == 2  # 1 synapse event (bias is not included)
+    assert summary_list_bias.summary_list[1].synapse_count == 2  # 2 synapses (single weight + bias)
+    assert summary_list_bias.summary_list[1].neuron_count == 0  # 0 neurons (layer is Linear (synaptic), no neurons)
+
+    assert summary_list_bias.summary_list[2].total_events == 1  # 1 synapse event (bias is not included)
+    assert summary_list_bias.summary_list[2].synapse_count == 0  # 0 synapses (layer is neuron, no synapses)
+    assert summary_list_bias.summary_list[
+               2].neuron_count == 1  # 1 neurons (layer is a neuron, output shape is 1, one event)
+
+    assert summary_list_bias.summary_list[3].total_events == 2  # 1 synapse event (bias is not included)
+    assert summary_list_bias.summary_list[3].synapse_count == 2  # 2 synapses (single weight + bias)
+    assert summary_list_bias.summary_list[3].neuron_count == 0  # 0 neurons (layer is Linear (synaptic), no neurons)
+
+    assert summary_list_bias.summary_list[4].total_events == 1  # 1 synapse event (bias is not included)
+    assert summary_list_bias.summary_list[4].synapse_count == 0  # 0 synapses (layer is neuron, no synapses)
+    assert summary_list_bias.summary_list[
+               4].neuron_count == 1  # 1 neurons (layer is a neuron, output shape is 1, one event)
+
+    # only one device specified, therefore length of energies should be 1, and equal to `expected_value_with_bias`,
+    # assert that with 1% tolerance
+    assert len(energies) == 1
+    assert pytest.approx(expected_value_with_bias, rel=0.01) == energies[0]
+
+
+def test_estimate_energy_network2():
+    """
+        More complicated architecture, testing the same things as test1 but
+
+            - non-neuromorphic architecture (cpu) with energy 8.6e-9 per synapse/neuron event
+
+            - 2 Layers architecture, much more events in comparison to test 1: ( not including bias in calculations )
+
+            self.fc1 = nn.Linear(32, 64)
+            self.lif1 = snn.Leaky(beta=beta, threshold=0.2)
+            self.fc2 = nn.Linear(64, 10)
+            self.lif2 = snn.Leaky(beta=beta, threshold=0.4)
+
+        Same code for keras (and nengo estimations)
+
+            # build an example model
+            inp = x = tf.keras.Input((1, 32))
+            x = tf.keras.layers.Dense(64)(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+            x = tf.keras.layers.Dense(10)(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+
+            model = tf.keras.Model(inp, x)
+            # estimate model energy
+            energy = keras_spiking.ModelEnergy(model)
+            energy.summary(print_warnings=False, columns=("name", "output_shape", "params", "connections", "neurons",
+                            "energy cpu-test3", "energy cpu"))
+
+            Layer (type)                     |Output shape |Param #|Conn #|Neuron #|J/inf (cpu-tes|J/inf (cpu)
+            ---------------------------------|-------------|-------|------|--------|--------------|-----------
+            input_3 (InputLayer)             |[(None, 1, 32|      0|     0|       0|             0|          0
+            dense_4 (Dense)                  |(None, 1, 64)|   2112|  2048|       0|         2e-05|    1.8e-05
+            spiking_activation_4 (SpikingActi|(None, 1, 64)|      0|     0|      64|       6.4e-08|    5.5e-07
+            dense_5 (Dense)                  |(None, 1, 10)|    650|   640|       0|       6.4e-06|    5.5e-06
+            spiking_activation_5 (SpikingActi|(None, 1, 10)|      0|     0|      10|         1e-08|    8.6e-08
+            ==================================================================================================
+            Total energy per inference [Joules/inf] (cpu-test3): 2.70e-05
+            Total energy per inference [Joules/inf] (cpu): 2.38e-05
+    """
+
+    # prepare the inputs
+    batch_size = 1
+    prob_of_one = 1
+    prob_matrix = prob_of_one * torch.ones((1, batch_size, 32))
+    inp_test = torch.bernoulli(prob_matrix)
+
+    # create the network (the "device under test")
+    energy_per_synapse_event = 8.6e-9
+    energy_per_neuron_event = 8.6e-9
+    DeviceProfileRegistry.add_device("cpu-test2", energy_per_synapse_event,
+                                     energy_per_neuron_event,
+                                     False)
+    v = EnergyEfficiencyNetTest2(beta=0.9)
+
+    # get the layer summaries
+    summary_list = estimate_energy(model=v, input_data=inp_test, devices="cpu-test2",
+                                   include_bias_term_in_events=False)
+
+    # nengo calculated value + as expected from simple calculations
+    expected_value = (2048 + 640) * energy_per_synapse_event + (64 + 10) * energy_per_neuron_event
+
+    energies = summary_list.total_energies
+    # calculate it for only one device
+    assert len(energies) == 1
+    assert pytest.approx(expected_value, rel=0.01) == energies[0], f"The energy estimate for simple network 2 " \
+                                                                   f"was not within tolerance (was {expected_value}, " \
+                                                                   f"is {energies[0]})"
+
+    #                                                   all of these numbers can be simplified, but written explicitly
+    #                                                   to show that they make sense (I think)
+    assert summary_list.summary_list[0].synapse_count == 2112 + 650
+    assert summary_list.summary_list[0].neuron_count == 64 + 10
+    assert summary_list.summary_list[0].total_events == (2112 - 64) + (650 - 10) + 64 + 10  # with bias removed
+
+    assert summary_list.summary_list[1].synapse_count == 2112
+    assert summary_list.summary_list[1].neuron_count == 0
+    assert summary_list.summary_list[1].total_events == 2112 - 64  # synapse count - bias events
+
+    assert summary_list.summary_list[2].synapse_count == 0
+    assert summary_list.summary_list[2].neuron_count == 64
+    assert summary_list.summary_list[2].total_events == 64
+
+    assert summary_list.summary_list[3].synapse_count == 650
+    assert summary_list.summary_list[3].neuron_count == 0
+    assert summary_list.summary_list[3].total_events == 650 - 10  # synapse count - bias events
+
+    assert summary_list.summary_list[4].synapse_count == 0
+    assert summary_list.summary_list[4].neuron_count == 10
+    assert summary_list.summary_list[4].total_events == 10
+
+    # now assert the same for when bias is included
+    summary_list = estimate_energy(model=v, input_data=inp_test, devices="cpu-test2",
+                                   include_bias_term_in_events=True)
+
+    # nengo calculated value + as expected from simple calculations (but now bias term is also included)
+    expected_value = (2112 + 650) * energy_per_synapse_event + (64 + 10) * energy_per_neuron_event
+
+    energies = summary_list.total_energies
+    # calculate it for only one device
+    assert len(energies) == 1
+    assert pytest.approx(expected_value, rel=0.01) == energies[0], f"The energy estimate for simple network 2 " \
+                                                                   f"was not within tolerance (was {expected_value}, " \
+                                                                   f"is {energies[0]})"
+
+    #                                                   all of these numbers can be simplified, but written explicitly
+    #                                                   to show that they make sense (I think)
+    assert summary_list.summary_list[0].synapse_count == 2112 + 650
+    assert summary_list.summary_list[0].neuron_count == 64 + 10
+    assert summary_list.summary_list[0].total_events == 2112 + 650 + 64 + 10
+
+    assert summary_list.summary_list[1].synapse_count == 2112
+    assert summary_list.summary_list[1].neuron_count == 0
+    assert summary_list.summary_list[1].total_events == 2112
+
+    assert summary_list.summary_list[2].synapse_count == 0
+    assert summary_list.summary_list[2].neuron_count == 64
+    assert summary_list.summary_list[2].total_events == 64
+
+    assert summary_list.summary_list[3].synapse_count == 650
+    assert summary_list.summary_list[3].neuron_count == 0
+    assert summary_list.summary_list[3].total_events == 650
+
+    assert summary_list.summary_list[4].synapse_count == 0
+    assert summary_list.summary_list[4].neuron_count == 10
+    assert summary_list.summary_list[4].total_events == 10
+
+
+def test_estimate_energy_network2_energies_for_two_new_devices():
+    """
+        Same architecture as above (EnergyEfficiencyNetTest2), but we create two different devices with two energy
+        values. We check whether the numbers reported for these devices are as expected
+
+            - non-neuromorphic architecture (cpu) with energy 8.6e-9 per synapse/neuron event
+            self.fc1 = nn.Linear(32, 64)
+            self.lif1 = snn.Leaky(beta=beta, threshold=0.2)
+            self.fc2 = nn.Linear(64, 10)
+            self.lif2 = snn.Leaky(beta=beta, threshold=0.4)
+
+        Using information from previous test / paper calculations :
+            Total synaptic events (without bias)  : 2048 + 640 = 2688
+            Total neuron events (without bias) : 64 + 10 = 74
+
+            Total synaptic events (with bias)  : 2112 + 650 = 2762
+            Total neuron events (with bias) : 64 + 10
+    """
+
+    # prepare the input
+    batch_size = 1
+    prob_of_one = 1
+    prob_matrix = prob_of_one * torch.ones((1, batch_size, 32))
+    inp_test = torch.bernoulli(prob_matrix)
+
+    # create the devices
+    device1_en_per_synapse, device1_en_per_neuron = 10e-9, 1e-9
+    device2_en_per_synapse, device2_en_per_neuron = 1e-9, 10e-9
+    DeviceProfileRegistry.add_device("cpu-test3", device1_en_per_synapse,
+                                     device1_en_per_neuron,
+                                     False)
+    DeviceProfileRegistry.add_device("cpu-test4", device2_en_per_synapse
+                                     , device2_en_per_neuron,
+                                     False)
+
+    # create the network (the "device under test")
+    v = EnergyEfficiencyNetTest2(beta=0.9)
+
+    # create the summary list
+    summary_list = estimate_energy(model=v, input_data=inp_test, devices=["cpu-test3", "cpu-test4"],
+                                   include_bias_term_in_events=False)
+    energies = summary_list.total_energies
+
+    # calculate it for two devices
+    expected_values = [None, None]
+    expected_values[0] = 2688 * device1_en_per_synapse + 74 * device1_en_per_neuron
+    expected_values[1] = 2688 * device2_en_per_synapse + 74 * device2_en_per_neuron
+
+    # check whether the returned energies are as expected
+    assert len(energies) == 2
+    assert pytest.approx(expected_values[0], rel=0.01) == energies[0]
+    assert pytest.approx(expected_values[1], rel=0.01) == energies[1]
+
+    # now does it but include bias
+
+    # create the summary list
+    summary_list = estimate_energy(model=v, input_data=inp_test, devices=["cpu-test3", "cpu-test4"],
+                                   include_bias_term_in_events=True)
+    energies = summary_list.total_energies
+
+    # calculate it for two devices
+    expected_values = [None, None]
+    expected_values[0] = 2762 * device1_en_per_synapse + 74 * device1_en_per_neuron
+    expected_values[1] = 2762 * device2_en_per_synapse + 74 * device2_en_per_neuron
+
+    # check whether the returned energies are as expected
+    assert len(energies) == 2
+    assert pytest.approx(expected_values[0], rel=0.01) == energies[0]
+    assert pytest.approx(expected_values[1], rel=0.01) == energies[1]
+
+
+def test_estimate_energy_network2_energies_for_batch_size_and_timesteps():
+    """
+        Check whether the energy/synapses/neuron counts scales linearly with batch_size
+        and check if the energy scales linearly with number of timesteps (for non-neuromorphic)
+    """
+
+    # create the devices
+    device1_en_per_synapse, device1_en_per_neuron = 1e-9, 1e-9
+    DeviceProfileRegistry.add_device("cpu-test-bs", device1_en_per_synapse,
+                                     device1_en_per_neuron,
+                                     False)
+
+    en = 2.762e-06
+    ts, bs = 1, 1
+    inp_test = torch.zeros((ts, bs, 32))
+
+    # create the network (the "device under test")
+    v = EnergyEfficiencyNetTest2(beta=0.9)
+
+    # create the summary list
+    summary_list = estimate_energy(model=v, input_data=inp_test, devices=["cpu-test-bs"],
+                                   include_bias_term_in_events=False)
+    data = summary_list.summary_list
+
+    # get the energy
+    base_counts = [(data[i].synapse_count, data[i].neuron_count, data[i].total_events) for i in range(0, 5)]
+    base_energy = summary_list.total_energies[0]
+
+    numbers_to_test = [1, 3, 5, 7, 11, 13]
+
+    for ts in numbers_to_test:
+        for bs in numbers_to_test:
+            inp_test = torch.zeros((ts, bs, 32))
+            net = EnergyEfficiencyNetTest2(beta=0.9)
+            summary_list = estimate_energy(model=net, input_data=inp_test, devices=["cpu-test-bs"],
+                                           include_bias_term_in_events=False)
+
+            # energy should scale linearly with timesteps (for non-neuromorphic) and batch_size (assert that)
+            assert pytest.approx(ts * bs * base_energy, rel=0.01) == summary_list.total_energies[0]
+
+            # total events should scale linearly with timesteps (for non-neuromorphic) and batch_size (assert that)
+            # synapse / neuron count should scale linearly with batch_size
+            for i in range(5):
+                # assert the synapse count
+                assert summary_list.summary_list[i].synapse_count == bs * base_counts[i][0]
+
+                # assert the neuron count
+                assert summary_list.summary_list[i].neuron_count == bs * base_counts[i][1]
+
+                # assert the total events count
+                assert summary_list.summary_list[i].total_events == ts * bs * base_counts[i][2]
+
+
+def test_estimate_energy_network3_for_neuromorphic():
+    """
+        First neuromorphic architecture test
+
+        Network is s modified EnergyEfficiencyNetTest1. It implements EnergyEstimationNetworkInterface, which makes it
+        easier to reset the network. It also implements "forward_full", which returns all the spikes, which makes it
+        slightly easier to test.
+
+        Input            : [0, 0, 0, 0, 0, 1, 0, 0, 0, 0]
+        1st Membrane     : [0, 0, 0, 0, 0, 1, 0.25, 0.0625, 0.03125]
+        1st spikes       : [0, 0, 0, 0, 0, 1, 0, 0, 0, 0]
+        2nd Membrane     : [0, 0, 0, 0, 0, 0.5, 0.125, 0.0625, 0.03125, 0.015625]
+        2nd spikes       : [0, 0, 0, 0, 0, 1, 0, 0, 0, 0]
+
+        There should be only 4 events : 2 synapses and 2 neuron activations
+    """
+
+    # prepare the input spike train
+    inp_test = torch.Tensor([[0, 0, 0, 0, 0, 1, 0, 0, 0, 0]]).T
+    # create the network (the "device under test")
+    v = EnergyEfficiencyNetTest3(beta=0.5)
+    v.fc1.weight.data = torch.Tensor([[1]])
+    v.fc1.bias.data = torch.Tensor([[0]])
+    v.fc2.weight.data = torch.Tensor([[0.5]])
+    v.fc2.bias.data = torch.Tensor([[0]])
+    v.reset()
+
+    # Debug, see whether the number checks out
+    DeviceProfileRegistry.add_device("neuromorphic-test1", 1e-9, 1e-9, True)
+    DeviceProfileRegistry.add_device("vn-against-neuromorphic-test1", 1e-9, 1e-9, False)
+
+    summary_list = estimate_energy(model=v, input_data=inp_test, devices=["neuromorphic-test1",
+                                                                          "vn-against-neuromorphic-test1"],
+                                   include_bias_term_in_events=False)
+
+    # see that the numbers differ a lot for devices that are truely spiking or not (by
+    # writing on paper, or checking on code, this set of weights will generate only one spike at time t=6)
+    # therefore expected energy is 4e-9 (1e-9 for each synapse and 1e-9 for each neuron)
+    # in comparison (in non spiking / von neunmann computing), although we have mostly zeros, energy will constantly
+    # be needed to perform operation
+    neuromorphic_expected = 4e-9
+    von_neunmann_expected = 4e-8
+    energies = summary_list.total_energies
+
+    # assert that the energies are as expected
+    assert len(energies) == 2
+    assert pytest.approx(neuromorphic_expected, rel=0.01) == energies[0]
+    assert pytest.approx(von_neunmann_expected, rel=0.01) == energies[1]
+
+    data = summary_list.summary_list
+
+    # assert the numbers for synapses/neurons/total events/spiking events
+    assert data[0].synapse_count == 4
+    assert data[0].neuron_count == 2
+    assert data[0].spiking_events == 4
+    assert data[0].total_events == 40
+
+    assert data[1].synapse_count == 2
+    assert data[1].neuron_count == 0
+    assert data[1].spiking_events == 1
+    assert data[1].total_events == 10
+
+    assert data[2].synapse_count == 0
+    assert data[2].neuron_count == 1
+    assert data[2].spiking_events == 1
+    assert data[2].total_events == 10
+
+    assert data[3].synapse_count == 2
+    assert data[3].neuron_count == 0
+    assert data[3].spiking_events == 1
+    assert data[3].total_events == 10
+
+
+def test_estimate_energy_network3_for_neuromorphic2():
+    """
+        Network is s modified EnergyEfficiencyNetTest1 (Simplest possible). It implements EnergyEstimationNetworkInterface,
+        which makes it easier to reset the network. It also implements "forward_full", which returns all the spikes,
+        which makes it slightly easier to test.
+
+        Input            : [1, 0, 0, 0, 0, 1, 0, 0, 0, 0]
+        1st spikes       : [1, 0, 0, 0, 0, 1, 0, 0, 0, 0]
+        2nd spikes       : [1, 0, 0, 0, 0, 1, 0, 0, 0, 0]
+
+        There should be only 4 events : 2 synapses and 2 neuron activations
+    """
+
+    """
+            Input spikes         :     1  0  0  0  1  0  0  0  0  0
+            Spiking activation 1 :     1  0  0  0  1  0  0  0  0  0
+            Spiking activation 2 :     1  0  0  0  1  0  0  0  0  0
+
+    """
+    # prepare the input
+    inp_test = torch.Tensor([[[1, 0, 0, 0, 1, 0, 0, 0, 0, 0]]]).T
+
+    # create the network (the "device under test")
+    v = EnergyEfficiencyNetTest3(beta=0)
+    v.fc1.weight.data = torch.Tensor([[1]])
+    v.fc1.bias.data = torch.Tensor([[0]])
+    v.fc2.weight.data = torch.Tensor([[1]])
+    v.fc2.bias.data = torch.Tensor([[0]])
+    v.reset()
+
+    DeviceProfileRegistry.add_device("neuromorphic-test2", 1e-9, 1e-9, True)
+    DeviceProfileRegistry.add_device("vn-against-neuromorphic-test2", 1e-9, 1e-9, False)
+
+    summary_list = estimate_energy(model=v, input_data=inp_test, devices=["neuromorphic-test2",
+                                                                          "vn-against-neuromorphic-test2"],
+                                   include_bias_term_in_events=False)
+
+    # see that the numbers differ a lot for devices that are truely spiking or not (by
+    # writing on paper, or checking on code, this set of weights will generate spikes at time t = 0 and t = 5)
+    # therefore expected energy is 8e-9 (1e-9 for each synapse and 1e-9 for each neuron)
+    neuromorphic_expected = (2 + 2 + 2 + 2) * 1e-9
+    von_neunmann_expected = 4 * 10 * 1e-9
+    energies = summary_list.total_energies
+
+    assert len(energies) == 2
+    assert pytest.approx(neuromorphic_expected, rel=0.01) == energies[0]
+    assert pytest.approx(von_neunmann_expected, rel=0.01) == energies[1]
+
+    # go through all the layers and assert the expected values
+    # assert the numbers for synapses/neurons/total events/spiking events
+    data = summary_list.summary_list
+    assert data[0].synapse_count == 4
+    assert data[0].neuron_count == 2
+    assert data[0].spiking_events == 8
+    assert data[0].total_events == 40
+
+    assert data[1].synapse_count == 2
+    assert data[1].neuron_count == 0
+    assert data[1].spiking_events == 2
+    assert data[1].total_events == 10
+
+    assert data[2].synapse_count == 0
+    assert data[2].neuron_count == 1
+    assert data[2].spiking_events == 2
+    assert data[2].total_events == 10
+
+    assert data[3].synapse_count == 2
+    assert data[3].neuron_count == 0
+    assert data[3].spiking_events == 2
+    assert data[3].total_events == 10
+
+
+def test_estimate_energy_network3_for_neuromorphic3():
+    """
+        Network is s modified EnergyEfficiencyNetTest1 (Simplest possible). It implements EnergyEstimationNetworkInterface,
+        which makes it easier to reset the network. It also implements "forward_full", which returns all the spikes,
+        which makes it slightly easier to test.
+
+            Input spikes         :     1  0  1  0  1  0  0  0  0  0
+            Spiking activation 1 :     1  0  1  1  1  0  0  0  0  0
+            Spiking activation 2 :     1  0  1  1  1  1  0  0  0  0
+    """
+
+    # prepare the input
+    inp_test = torch.Tensor([[1, 0, 1, 0, 1, 0, 0, 0, 0, 0]]).T
+
+    # create the network (the "device under test")
+    v = EnergyEfficiencyNetTest3(beta=0.75)
+    v.fc1.weight.data = torch.Tensor([[1]])
+    v.fc1.bias.data = torch.Tensor([[0]])
+    v.fc2.weight.data = torch.Tensor([[0.5]])
+    v.fc2.bias.data = torch.Tensor([[0]])
+    v.reset()
+
+    DeviceProfileRegistry.add_device("neuromorphic-test3", 1e-9, 1e-9, True)
+    DeviceProfileRegistry.add_device("vn-against-neuromorphic-test3", 1e-9, 1e-9, False)
+
+    summary_list = estimate_energy(model=v, input_data=inp_test, devices=["neuromorphic-test3",
+                                                                          "vn-against-neuromorphic-test3"],
+                                   include_bias_term_in_events=False)
+
+    # on input 3 spikes, in layer 1 we have 4 spikes, in final layer 5 spikes
+    # - 3 synapse events (on layer 1)
+    # - 4 neuron events in layer 1 activation
+    # - 4 synapse events in layer 2
+    # - 5 neuron events in layer 2 activation
+    neuromorphic_expected = (3 + 4 + 4 + 5) * 1e-9
+    von_neunmann_expected = 4 * 10 * 1e-9
+    energies = summary_list.total_energies
+
+    # check that the calculated energies are the same as expected
+    assert len(energies) == 2
+    assert pytest.approx(neuromorphic_expected, rel=0.01) == energies[0]
+    assert pytest.approx(von_neunmann_expected, rel=0.01) == energies[1]
+
+    # go through all the layers and assert the expected values
+    # assert the numbers for synapses/neurons/total events/spiking events
+    data = summary_list.summary_list
+    assert data[0].spiking_events == 3 + 4 + 4 + 5
+    assert data[0].total_events == 40
+
+    assert data[1].spiking_events == 3
+    assert data[1].total_events == 10
+
+    assert data[2].spiking_events == 4
+    assert data[2].total_events == 10
+
+    assert data[3].spiking_events == 4
+    assert data[3].total_events == 10
+
+    assert data[4].spiking_events == 5
+    assert data[4].total_events == 10
+
+
+def test_estimate_energy_network4_for_neuromorphic4():
+    """
+            CNN test
+            Input spikes         :     p1  p2
+
+            pattern 1 (p1) = [[1 0 0 0 1]
+                              [0 1 0 1 0]
+                              [0 0 1 0 0]
+                              [0 1 0 1 0]
+                              [1 0 0 0 1]]
+
+            pattern 2 (p2) = [[0 1 1 1 0]
+                              [1 0 0 0 1]
+                              [1 0 0 0 1]
+                              [1 0 0 0 1]
+                              [0 1 1 1 0]]
+
+
+            CNN spike counts for the CNN layer:
+            p1 : [[3 3 3]
+                  [3 5 3]
+                  [3 3 3]]
+
+            total spiking events (sum ) : 29
+
+            p2 :  [[4 3 4]
+                   [3 0 3]
+                   [4 3 4]]
+
+            total spiking events (sum) : 28
+
+            expected total spiking events in 2D CNN = 29 + 28 = 57
+
+
+            Total Events (excluding the bias event) :
+            - map produces 3 by 3 image (sliding a kernel of size 3x3) : 9 * 9 = 81 events
+            - neuron layer : 3 x 3 image, each pixel one event = 9 events
+            - dense layer : flattens the input map to size [1, 9] = 9 events
+            - neuron layer : 1 event
+
+            2 timesteps, therefore = 2 * (81 + 9 + 9 + 1) = 200 events
+    """
+    p1 = [[1, 0, 0, 0, 1],
+          [0, 1, 0, 1, 0],
+          [0, 0, 1, 0, 0],
+          [0, 1, 0, 1, 0],
+          [1, 0, 0, 0, 1]]
+
+    p2 = [[0, 1, 1, 1, 0],
+          [1, 0, 0, 0, 1],
+          [1, 0, 0, 0, 1],
+          [1, 0, 0, 0, 1],
+          [0, 1, 1, 1, 0]]
+
+    inp_test = torch.Tensor([[p1], [p2]])
+    v = EnergyEfficiencyNetTest4()
+    v.reset()
+    v.cnn1.weight.data = torch.Tensor([[[[0.75, 0, 0],
+                                         [0, 0.75, 0],
+                                         [0, 0, 0.75]]]])
+    v.cnn1.bias.data = torch.Tensor([0])
+    v.fc2.weight.data = torch.Tensor([[0, 0, 0, 0, 2, 0, 0, 0, 0]])
+    v.fc2.bias.data = torch.Tensor([0])
+
+    DeviceProfileRegistry.add_device("neuromorphic-test8", 1e-9, 1e-9, True)
+    DeviceProfileRegistry.add_device("vn-against-neuromorphic-test8", 1e-9, 1e-9, False)
+    summary_list = estimate_energy(model=v, input_data=inp_test, include_bias_term_in_events=False,
+                                   network_requires_first_dim_as_time=False)
+
+    # check whether the number of total/spiking events is correct for root layer
+    assert summary_list.summary_list[0].spiking_events == 68
+    assert summary_list.summary_list[0].total_events == 200
+
+    # check whether the number of total/spiking events is correct for first CNN layer
+    assert summary_list.summary_list[1].spiking_events == (29 + 28)
+    assert summary_list.summary_list[1].total_events == 2 * (9 * 9)
+
+    # check whether the number of total/spiking events is correct for first neuron layer
+    assert summary_list.summary_list[2].spiking_events == 5
+    assert summary_list.summary_list[2].total_events == 2 * 9
+
+    # check whether the number of total/spiking events is correct for dense layer
+    assert summary_list.summary_list[3].spiking_events == 5
+    assert summary_list.summary_list[3].total_events == 2 * 9
+
+    # check whether the total/spiking events is correct for classification layer
+    assert summary_list.summary_list[4].spiking_events == 1
+    assert summary_list.summary_list[4].total_events == 2
+
+
+def test_network_large():
+    """
+        Compare a large network to equivalent nengo keras_spiking network.
+        Below is the keras equivalent of pytorch network. make sure that the energy estimates and total
+        event counts are the same in comparison to nengo
+
+            # build the equivalent model
+            inp = x = tf.keras.Input((1, 32, 32, 1))
+            x = tf.keras.layers.Conv2D(16, kernel_size=3, padding="same")(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+            x = tf.keras.layers.Conv2D(16, kernel_size=5, dilation_rate=2, padding="valid")(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+            x = tf.keras.layers.Conv2D(16, kernel_size=7, dilation_rate=3, padding="valid")(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+            dense = tf.keras.layers.Reshape((1, 6 * 6 * 16))(x)
+            dense = tf.keras.layers.Dense(32)(dense)
+            dense = keras_spiking.SpikingActivation("relu")(dense)
+            dense = tf.keras.layers.Dense(10)(dense)
+            dense = keras_spiking.SpikingActivation("relu")(dense)
+            model = tf.keras.Model(inp, dense)
+
+            # estimate model energy
+            energy = keras_spiking.ModelEnergy(model)
+            energy.summary(print_warnings=False, columns=("name", "output_shape", "params", "connections", "neurons", "energy cpu"))
+
+        Output:
+            Layer (type)                           |Output shape         |Param #|Conn # |Neuron #|J/inf (cpu)
+            ---------------------------------------|---------------------|-------|-------|--------|-----------
+            input_1 (InputLayer)                   |[(None, 1, 32, 32, 1)|      0|      0|       0|          0
+            conv2d (Conv2D)                        |(None, 1, 32, 32, 16)|    160| 147456|       0|     0.0013
+            spiking_activation (SpikingActivation) |(None, 1, 32, 32, 16)|      0|      0|   16384|    0.00014
+            conv2d_1 (Conv2D)                      |(None, 1, 24, 24, 16)|   6416|3686400|       0|      0.032
+            spiking_activation_1 (SpikingActivation|(None, 1, 24, 24, 16)|      0|      0|    9216|    7.9e-05
+            conv2d_2 (Conv2D)                      |  (None, 1, 6, 6, 16)|  12560| 451584|       0|     0.0039
+            spiking_activation_2 (SpikingActivation|  (None, 1, 6, 6, 16)|      0|      0|     576|      5e-06
+            reshape (Reshape)                      |       (None, 1, 576)|      0|      0|       0|          0
+            dense (Dense)                          |        (None, 1, 32)|  18464|  18432|       0|    0.00016
+            spiking_activation_3 (SpikingActivation|        (None, 1, 32)|      0|      0|      32|    2.8e-07
+            dense_1 (Dense)                        |        (None, 1, 10)|    330|    320|       0|    2.8e-06
+            spiking_activation_4 (SpikingActivation|        (None, 1, 10)|      0|      0|      10|    8.6e-08
+            ==================================================================================================
+            Total energy per inference [Joules/inf] (cpu): 3.72e-02
+
+    """
+    p1 = np.ones((1, 32, 32))
+
+    inp_test = torch.Tensor(np.array([[p1]]))
+    v = EnergyEfficiencyNetTest5Large()
+    v.reset()
+
+    # it's important to set the network_requires_first_dim_as_time to False : the 2D convolutions does not accept the
+    # shape of (time, batch, W, H, C), therefore we want to use the input as (batch, W, H, C), which is done by setting
+    # network_requires_first_dim_as_time = False
+    summary_list = estimate_energy(model=v, input_data=inp_test, include_bias_term_in_events=False,
+                                   network_requires_first_dim_as_time=False)
+
+    # assert total counts
+    assert summary_list.summary_list[0].total_events == (147456 + 16384 + 3686400 + 9216 + 451584 + 576 + 18432 + 32 +
+                                                         320 + 10)
+
+    # assert the total count for first convolutional layer
+    assert summary_list.summary_list[1].total_events == 147456
+
+    # assert the total count for first spiking layer
+    assert summary_list.summary_list[2].total_events == 16384
+
+    # assert the total count for second convolutional layer
+    assert summary_list.summary_list[3].total_events == 3686400
+
+    # assert the total count for second spiking layer
+    assert summary_list.summary_list[4].total_events == 9216
+
+    # assert the total count for third convolutional layer
+    assert summary_list.summary_list[5].total_events == 451584
+
+    # assert the total count for third spiking layer
+    assert summary_list.summary_list[6].total_events == 576
+
+    # assert the total count for first dense layer
+    assert summary_list.summary_list[7].total_events == 18432
+
+    # assert the total count for fourth neuron spiking layer
+    assert summary_list.summary_list[8].total_events == 32
+
+    # assert the total count for second dense layer (classification head)
+    assert summary_list.summary_list[9].total_events == 320
+
+    # assert the total count for fifth neuron spiking layer ( classification head)
+    assert summary_list.summary_list[10].total_events == 10
+
+    # assert that the total energy calculated by this new module is roughly the same as calculated in nengo
+    energies = summary_list.total_energies
+    assert len(energies) == 1
+
+    assert pytest.approx(3.72e-02, rel=0.01) == energies[0]
+
+
+def test_network_1D_conv():
+    """
+        Compare a large network to equivalent nengo keras_spiking network, do it now for 1D convolutions.
+        Below is the keras equivalent of pytorch network. make sure that the energy estimates and total
+        event counts are the same in comparison to nengo
+
+            # build an example model
+            inp = x = tf.keras.Input((1, 32, 1))
+            x = tf.keras.layers.Conv1D(16, kernel_size=3, padding="same")(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+            x = tf.keras.layers.Conv1D(16, kernel_size=5, padding="valid")(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+            dense = tf.keras.layers.Reshape((1, 16 * 28))(x)
+            dense = tf.keras.layers.Dense(64)(dense)
+            dense = keras_spiking.SpikingActivation("relu")(dense)
+            dense = tf.keras.layers.Dense(1)(dense)
+            dense = keras_spiking.SpikingActivation("relu")(dense)
+
+            model = tf.keras.Model(inp, dense)
+
+            # estimate model energy
+            energy = keras_spiking.ModelEnergy(model)
+            energy.summary(print_warnings=False, columns=("name", "output_shape", "params", "connections", "neurons", "energy cpu"))
+
+        Output:
+             Layer (type)                            |Output shape      |Param #|Conn #|Neuron #|J/inf (cpu)
+            ----------------------------------------|------------------|-------|------|--------|-----------
+            input_3 (InputLayer)                    |[(None, 1, 32, 1)]|      0|     0|       0|          0
+            conv1d_1 (Conv1D)                       | (None, 1, 32, 16)|     64|  1536|       0|    1.3e-05
+            spiking_activation_6 (SpikingActivation)| (None, 1, 32, 16)|      0|     0|     512|    4.4e-06
+            conv1d_2 (Conv1D)                       | (None, 1, 28, 16)|   1296| 35840|       0|    0.00031
+            spiking_activation_7 (SpikingActivation)| (None, 1, 28, 16)|      0|     0|     448|    3.9e-06
+            reshape_1 (Reshape)                     |    (None, 1, 448)|      0|     0|       0|          0
+            dense_2 (Dense)                         |     (None, 1, 64)|  28736| 28672|       0|    0.00025
+            spiking_activation_8 (SpikingActivation)|     (None, 1, 64)|      0|     0|      64|    5.5e-07
+            dense_3 (Dense)                         |      (None, 1, 1)|     65|    64|       0|    5.5e-07
+            spiking_activation_9 (SpikingActivation)|      (None, 1, 1)|      0|     0|       1|    8.6e-09
+            ===============================================================================================
+            Total energy per inference [Joules/inf] (cpu): 5.77e-04
+    """
+    p1 = np.ones((1, 32))
+
+    inp_test = torch.Tensor(np.array([p1]))
+    v = EnergyEfficiencyNet1DConv()
+    v.reset()
+
+    summary_list = estimate_energy(model=v, input_data=inp_test, include_bias_term_in_events=False,
+                                   network_requires_first_dim_as_time=False)
+
+    # assert total counts
+    assert summary_list.summary_list[0].total_events == (1536 + 512 + 35840 + 448 + 28672 + 64 + 64 + 1)
+
+    # assert the total count for first convolutional layer
+    assert summary_list.summary_list[1].total_events == 1536
+
+    # assert the total count for first spiking layer
+    assert summary_list.summary_list[2].total_events == 512
+
+    # assert the total count for second convolutional layer
+    assert summary_list.summary_list[3].total_events == 35840
+
+    # assert the total count for second spiking layer
+    assert summary_list.summary_list[4].total_events == 448
+
+    # assert the total count for first dense layer
+    assert summary_list.summary_list[5].total_events == 28672
+
+    # assert the total count for third spiking layer
+    assert summary_list.summary_list[6].total_events == 64
+
+    # assert the total count for second dense layer (classification)
+    assert summary_list.summary_list[7].total_events == 64
+
+    # assert the total count for fourth neuron spiking layer ( classification)
+    assert summary_list.summary_list[8].total_events == 1
+
+    # assert that the total energy calculated by this new module is roughly the same as calculated in nengo
+    energies = summary_list.total_energies
+    assert len(energies) == 1
+
+    assert pytest.approx(5.77e-04, rel=0.01) == energies[0]
+
+
+def test_network_unrecognized_leaf_layers():
+    """
+        Test how the network handles the unrecognized layers : for these, it should set the value to None, which will
+        be displayed as '?' in the print command.
+
+        Later test registering this new layer (which is a Linear layer underneath)
+    """
+    p1 = np.ones((1, 128))
+
+    inp_test = torch.Tensor(np.array([p1]))
+    v = EnergyEfficiencyNetUnrecognizedLayers()
+    v.reset()
+
+    # do a forward pass, where we have a "UnrecognizedLeafTestLayer" (which under the hood is linear layer)
+    # the values for these layers should be unrecognized (None). Assert that these values are None
+    summary_list = estimate_energy(model=v, input_data=inp_test, include_bias_term_in_events=False)
+
+    assert summary_list.summary_list[0].neuron_count == 32 + 16
+    # TODO : should None here be propagated or not ?
+    assert summary_list.summary_list[0].synapse_count == 0
+
+    # First layer (the unrecognized).Doesn't have a callbacks, therefore synapse, neuron, total, spiking values are None
+    assert summary_list.summary_list[1].neuron_count is None
+    assert summary_list.summary_list[1].synapse_count is None
+    assert summary_list.summary_list[1].total_events is None
+    assert summary_list.summary_list[1].spiking_events is None
+
+    # second layer (SNN spiking neurons). Here expect the neuron count of 32
+    assert summary_list.summary_list[2].neuron_count == 32
+    assert summary_list.summary_list[2].synapse_count == 0
+    assert summary_list.summary_list[2].total_events == 32
+
+    # third layer (the unrecognized).Doesn't have a callbacks, therefore synapse, neuron, total, spiking values
+    # are None
+    assert summary_list.summary_list[3].neuron_count is None
+    assert summary_list.summary_list[3].synapse_count is None
+    assert summary_list.summary_list[3].total_events is None
+    assert summary_list.summary_list[3].spiking_events is None
+
+    # fourth layer (SNN spiking neurons). Here expect the neuron count of 16
+    assert summary_list.summary_list[4].neuron_count == 16
+    assert summary_list.summary_list[4].synapse_count == 0
+    assert summary_list.summary_list[4].total_events == 16
+
+    # now add the device to registry with the callback to the same function which handles the linear, and assert
+    # as expected
+    LayerParameterEventCalculator.register_new_layer(UnrecognizedLeafTestLayer,
+                                                     synapse_neuron_count_for_linear,
+                                                     count_events_for_linear)
+    # do a forward pass, where we have a "UnrecognizedLeafTestLayer" (which under the hood is linear layer)
+    # now the layers should be recognized, and no unknown values (None) should be there
+    summary_list = estimate_energy(model=v, input_data=inp_test, include_bias_term_in_events=False)
+
+    assert summary_list.summary_list[0].neuron_count == 32 + 16
+    # TODO : should None here be propagated or not ?
+    assert summary_list.summary_list[0].synapse_count == 4128 + 528
+
+    # First layer (the unrecognized).Doesn't have a callbacks, therefore synapse, neuron, total, spiking values are None
+    assert summary_list.summary_list[1].neuron_count == 0
+    assert summary_list.summary_list[1].synapse_count == 4128
+    assert summary_list.summary_list[1].total_events == 4096
+    assert summary_list.summary_list[1].spiking_events is not None
+
+    # second layer (SNN spiking neurons). Here expect the neuron count of 32
+    assert summary_list.summary_list[2].neuron_count == 32
+    assert summary_list.summary_list[2].synapse_count == 0
+    assert summary_list.summary_list[2].total_events == 32
+
+    # third layer (the unrecognized).Doesn't have a callbacks, therefore synapse, neuron, total, spiking values
+    # are None
+    assert summary_list.summary_list[3].neuron_count == 0
+    assert summary_list.summary_list[3].synapse_count == 528
+    assert summary_list.summary_list[3].total_events == 512
+    assert summary_list.summary_list[3].spiking_events is not None
+
+    # fourth layer (SNN spiking neurons). Here expect the neuron count of 16
+    assert summary_list.summary_list[4].neuron_count == 16
+    assert summary_list.summary_list[4].synapse_count == 0
+    assert summary_list.summary_list[4].total_events == 16
+
+
+def test_nested_network():
+    """
+        test whether the numbers are as expected for a pytorch network consisting of few, nested nn.modules.
+        We set the depth to 3, and get all the information about all nn.modules (from the topmost layer to basic Linear
+        layer). We set the synapse/neuron/total event count to expected values (as in unfolded network)
+
+        # unfolded pytorch model
+        class EnergyEfficiencyNestedNetworkTestUnfolded(EnergyEstimationNetworkInterface):
+            def __init__(self, Beta=0.9):
+                super().__init__()
+                self.fc1 = nn.Linear(128, 64)
+                self.lif1 = snn.Leaky(Beta)
+
+                self.fc2 = nn.Linear(64, 32)
+                self.lif2 = snn.Leaky(Beta)
+
+                self.fc3 = nn.Linear(32, 16)
+                self.lif3 = snn.Leaky(Beta)
+
+                self.fc4 = nn.Linear(16, 32)
+                self.lif4 = snn.Leaky(Beta)
+
+                self.fc5 = nn.Linear(32, 64)
+                self.lif5 = snn.Leaky(Beta)
+
+                self.fc6 = nn.Linear(64, 128)
+                self.lif6 = snn.Leaky(Beta)
+                self.reset()
+
+            def forward(self, x: torch.Tensor):
+                spk, self.mem1 = self.lif1(self.fc1(x), self.mem1)
+                spk, self.mem2 = self.lif2(self.fc2(spk), self.mem2)
+                spk, self.mem3 = self.lif3(self.fc3(spk), self.mem3)
+                spk, self.mem4 = self.lif4(self.fc4(spk), self.mem4)
+                spk, self.mem5 = self.lif5(self.fc5(spk), self.mem5)
+                return self.lif6(self.fc6(spk), self.mem6)
+
+    -------------------------------------------------------------------------------------
+        # Keras/nengo unfolded equivalent model
+            # build an example model
+            inp = x = tf.keras.Input((1, 128))
+            x = tf.keras.layers.Dense(64)(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+
+            x = tf.keras.layers.Dense(32)(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+
+            x = tf.keras.layers.Dense(16)(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+
+            x = tf.keras.layers.Dense(32)(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+
+            x = tf.keras.layers.Dense(64)(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+
+            x = tf.keras.layers.Dense(128)(x)
+            x = keras_spiking.SpikingActivation("relu")(x)
+
+
+            model = tf.keras.Model(inp, x)
+
+            # estimate model energy
+            energy = keras_spiking.ModelEnergy(model)
+            energy.summary(print_warnings=False, columns=("name", "output_shape", "params", "connections", "neurons", "energy cpu"))
+
+
+        Layers are as follows :
+        0 Model
+        1     Block
+        2         Module
+        3             Linear
+        4             Leaky
+        5         Module
+        6             Linear
+        7             Leaky
+        8         Module
+        9             Linear
+        10            Leaky
+        11    Block
+        12        Module
+        13            Linear
+        14            Leaky
+        15        Module
+        16            Linear
+        17            Leaky
+        18        Module
+        19            Linear
+        20            Leaky
+
+    """
+
+
+    p1 = np.ones((1, 128))
+
+    inp_test = torch.Tensor(np.array([p1]))
+    v = EnergyEfficiencyNestedNetworkTest()
+    v.reset()
+
+    n = EnergyEfficiencyNestedNetworkTestUnfolded()
+    synapses = [8256, 2080, 528, 544, 2112, 8320]
+    neurons = [64, 32, 16, 32, 64, 128]
+
+    energy_per_neuron, energy_per_synapse = 1e-9, 1e-9
+    predicted_energy = (sum(synapses) - sum(neurons)) * energy_per_synapse + sum(neurons) * energy_per_neuron
+    DeviceProfileRegistry.add_device("nested-network-test", energy_per_synapse,
+                                     energy_per_neuron, False)
+
+    # do a forward pass, where we have a "UnrecognizedLeafTestLayer" (which under the hood is linear layer)
+    # the values for these layers should be unrecognized (None). Assert that these values are None
+    summary_list = estimate_energy(model=v, input_data=inp_test, include_bias_term_in_events=False,
+                                   depth=3, devices="nested-network-test")
+
+    data = summary_list.summary_list
+
+    # assert that the energy is as predicted
+    assert pytest.approx(predicted_energy, rel=0.01) == summary_list.total_energies[0]
+
+    assert data[0].synapse_count == sum(synapses)
+    assert data[0].neuron_count == sum(neurons)
+    assert data[0].total_events == sum(synapses) # this is valid when bias is not included, however when bias is included
+                                                 # it should be sum(synapses) + sum(neurons)
+
+    # first block (module -> module -> module)
+    assert data[1].synapse_count == sum(synapses[:3])
+    assert data[1].neuron_count == sum(neurons[:3])
+    assert data[1].total_events == sum(synapses[:3])
+
+    # second block (module -> module -> module)
+    assert data[11].synapse_count == sum(synapses[3:])
+    assert data[11].neuron_count == sum(neurons[3:])
+    assert data[11].total_events == sum(synapses[3:])
+
+    for offset, counts_offset in zip([1, 11], [0, 3]):
+        for i in range(3):
+            # check the modules (Linear -> Leaky)
+            count_idx = counts_offset + i
+            assert data[1 + offset + 3*i].synapse_count == synapses[count_idx]
+            assert data[1 + offset + 3*i].neuron_count == neurons[count_idx]
+            assert data[1 + offset + 3*i].total_events == synapses[count_idx]
+
+            # check the individual layers
+            assert data[1 + offset + 3 * i + 1].synapse_count == synapses[count_idx]
+            assert data[1 + offset + 3 * i + 1].neuron_count == 0
+            assert data[1 + offset + 3 * i + 1].total_events == synapses[count_idx] - neurons[count_idx]
+
+            assert data[1 + offset + 3 * i + 2].synapse_count == 0
+            assert data[1 + offset + 3 * i + 2].neuron_count == neurons[count_idx]
+            assert data[1 + offset + 3 * i + 2].total_events == neurons[count_idx]
+
+
+def test_nested_network_depth1():
+    """
+        test whether the numbers are as expected for a pytorch network consisting of few, nested nn.modules.
+        We set the depth to 1, and get only information about two blocks (however the number of parameters should be
+        the same for these blocks as in test above)
+
+    """
+
+
+    p1 = np.ones((1, 128))
+
+    inp_test = torch.Tensor(np.array([p1]))
+    v = EnergyEfficiencyNestedNetworkTest()
+    v.reset()
+
+    n = EnergyEfficiencyNestedNetworkTestUnfolded()
+    synapses = [8256, 2080, 528, 544, 2112, 8320]
+    neurons = [64, 32, 16, 32, 64, 128]
+
+    energy_per_neuron, energy_per_synapse = 1e-9, 1e-9
+    predicted_energy = (sum(synapses) - sum(neurons)) * energy_per_synapse + sum(neurons) * energy_per_neuron
+    DeviceProfileRegistry.add_device("nested-network-test", energy_per_synapse,
+                                     energy_per_neuron, False)
+
+    # do a forward pass, where we have a "UnrecognizedLeafTestLayer" (which under the hood is linear layer)
+    # the values for these layers should be unrecognized (None). Assert that these values are None
+    summary_list = estimate_energy(model=v, input_data=inp_test, include_bias_term_in_events=False,
+                                   depth=1, devices="nested-network-test")
+
+    data = summary_list.summary_list
+
+    # assert that the energy is as predicted
+    assert pytest.approx(predicted_energy, rel=0.01) == summary_list.total_energies[0]
+
+    assert data[0].synapse_count == sum(synapses)
+    assert data[0].neuron_count == sum(neurons)
+    assert data[0].total_events == sum(synapses) # this is valid when bias is not included, however when bias is included
+                                                 # it should be sum(synapses) + sum(neurons)
+
+    # first block (module -> module -> module)
+    assert data[1].synapse_count == sum(synapses[:3])
+    assert data[1].neuron_count == sum(neurons[:3])
+    assert data[1].total_events == sum(synapses[:3])
+
+    # second block (module -> module -> module)
+    assert data[11].synapse_count == sum(synapses[3:])
+    assert data[11].neuron_count == sum(neurons[3:])
+    assert data[11].total_events == sum(synapses[3:])