"""
Contains the module base classes for Rockpool
"""
# - Rockpool imports
import collections
from rockpool.parameters import ParameterBase
from rockpool.timeseries import TimeSeries
try:
from rockpool.graph.graph_base import GraphModuleBase
except:
GraphModuleBase = "GraphModuleBase"
# - Other imports
from abc import ABC, abstractmethod
from warnings import warn
from collections import ChainMap
from typing import Tuple, Any, Iterable, Dict, Optional, List, Union
import numpy as np
__all__ = ["Module"]
class ModuleBase(ABC):
"""
Base class for all `Module` subclasses in Rockpool
"""
def __init__(
self,
shape: Optional[Union[Tuple, int]] = None,
spiking_input: bool = False,
spiking_output: bool = False,
*args,
**kwargs,
):
"""
Initialise this module
Args:
shape (Optional[Union[Tuple, int]]): The shape of the defined module
spiking_input (bool): Whether this module receives spiking input. Default: False
spiking_output (bool): Whether this module produces spiking output. Default: False
*args: Additional positional arguments
**kwargs: Additional keyword arguments
"""
# - Set flag to specify that we are in the `__init__()` method
self._in_Module_init = True
""" (bool) If exists and ``True``, indicates that the module is in the ``__init__`` chain."""
self._force_set_attributes = False
""" (bool) If ``True``, do not sanity-check attributes when setting. """
# - Initialise co-classes etc.
super().__init__(*args, **kwargs)
# - Initialise Module attributes
self._submodulenames: List[str] = []
"""Registry of sub-module names"""
self._name: Optional[str] = None
"""Name of this module, if assigned"""
self._spiking_input: bool = spiking_input
"""Whether this module receives spiking input"""
self._spiking_output: bool = spiking_output
"""Whether this module produces spiking output"""
# - Be generous if a scalar was provided instead of a tuple
if isinstance(shape, Iterable):
self._shape = tuple(shape)
"""The shape of this module"""
else:
self._shape = (shape,)
"""The shape of this module"""
def __repr__(self, indent: str = "") -> str:
"""
Produce a string representation of this module
Args:
indent (str): The indent to prepend to each line of output
Returns:
str: A string representation of this module
"""
# - String representation
repr = f"{indent}{self.full_name} with shape {self._shape}"
# - Add submodules
if self._submodulenames:
repr += " {"
for mod_name in self._submodulenames:
repr += "\n" + ModuleBase.__repr__(
getattr(self, mod_name), indent=indent + " "
)
repr += f"\n{indent}" + "}"
return repr
def _get_attribute_registry(self) -> Tuple[Dict, Dict]:
"""
Return or initialise the attribute registry for this module
Returns:
(tuple): registered_attributes, registered_modules
"""
if not hasattr(self, "_ModuleBase__registered_attributes") or not hasattr(
self, "_ModuleBase__modules"
):
super().__setattr__("_ModuleBase__registered_attributes", {})
super().__setattr__("_ModuleBase__modules", {})
# - Get the attribute and modules dictionaries in a safe way
__registered_attributes = self.__dict__.get(
"_ModuleBase__registered_attributes"
)
__modules = self.__dict__.get("_ModuleBase__modules")
return __registered_attributes, __modules
def __setattr__(self, name: str, val: Any):
"""
Set an attribute for this module
Args:
name (str): The name of the attribute to set
val (Any): The value to assign to the attribute
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
# - Check if this is a new rockpool Parameter
if isinstance(val, ParameterBase):
try:
super().__getattribute__("_in_Module_init")
except Exception as e:
raise NotImplementedError(
"You must call `super.__init__()` in your `Module` subclass."
)
if (
hasattr(self, name)
and hasattr(self, "_in_Module_init")
and not self._in_Module_init
):
raise ValueError(
f'Cannot assign a new Parameter or State to an existing attribute "{name}".'
)
# - Register the attribute
self._register_attribute(name, val)
val = val.data
# - Are we assigning a sub-module?
if isinstance(val, ModuleBase):
self._register_module(name, val)
# - Check if this is an already registered attribute
if name in __registered_attributes:
if hasattr(self, name):
(_, _, _, _, shape) = __registered_attributes[name]
if val is not None:
# - Should we force-set the attribute?
if self._force_set_attributes:
__registered_attributes[name][4] = np.shape(val)
elif np.shape(val) != shape:
# - Check that shapes are identical
raise ValueError(
f"The new value assigned to {name} must be of shape {shape} (got {np.shape(val)})."
)
# - Assign the value to the __registered_attributes dictionary
__registered_attributes[name][0] = val
# - Assign attribute to self
super().__setattr__(name, val)
def __delattr__(self, name: str):
"""
Delete an attribute from this module, and remove from the attribute registry if present
Args:
name (str): The name of the attribute to delete
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
# - Remove attribute from registered attributes
if name in __registered_attributes:
del __registered_attributes[name]
# - Remove name from modules
if name in __modules:
del __modules[name]
self._submodulenames.remove(name)
# - Remove attribute
super().__delattr__(name)
def _register_attribute(self, name: str, val: ParameterBase):
"""
Record an attribute in the attribute registry
Args:
name (str): The name of the attribute to register
val (ParameterBase): The `ParameterBase` subclass object to register. e.g. `Parameter`, `SimulationParameter` or `State`.
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
# - Record attribute in attribute registry
__registered_attributes[name]: dict = [
val.data,
type(val).__name__,
val.family,
val.init_func,
val.shape,
]
"""The attribute registry for this module"""
def _register_module(self, name: str, mod: "ModuleBase"):
"""
Register a sub-module in the module registry
Args:
name (str): The name of the module to register
mod (ModuleBase): The `ModuleBase` object to register
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
if not isinstance(mod, ModuleBase):
raise ValueError(
f"You may only assign a `Module` subclass as a sub-module."
)
# - Assign module name to module
mod._name = name
# - Assign to appropriate attribute dictionary
__modules[name] = [mod, type(mod).__name__]
if name not in self._submodulenames:
self._submodulenames.append(name)
def set_attributes(self, new_attributes: dict) -> "ModuleBase":
"""
Set the attributes and sub-module attributes from a dictionary
This method can be used with the dictionary returned from module evolution to set the new state of the module. It can also be used to set multiple parameters of a module and submodules.
Examples:
Use the functional API to evolve, obtain new states, and set those states:
>>> _, new_state, _ = mod(input)
>>> mod = mod.set_attributes(new_state)
Obtain a parameter dictionary, modify it, then set the parameters back:
>>> params = mod.parameters()
>>> params['w_input'] *= 0.
>>> mod.set_attributes(params)
Args:
new_attributes (dict): A nested dictionary containing parameters of this module and sub-modules.
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
# - Set self attributes
for k, v in __registered_attributes.items():
if k in new_attributes:
self.__setattr__(k, new_attributes[k])
# - Set submodule attributes
for k, m in __modules.items():
if k in new_attributes:
m[0].set_attributes(new_attributes[k])
# - Return the module, for compatibility with the functional interface
return self
def _get_attribute_family(
self, type_name: str, family: Union[str, Tuple, List] = None
) -> dict:
"""
Search for attributes of this module and submodules that match a given family
This method can be used to conveniently get all weights for a network; or all time constants; or any other family of parameters. Parameter families are defined simply by a string: ``"weights"`` for weights; ``"taus"`` for time constants, etc. These strings are arbitrary, but if you follow the conventions then future developers will thank you (that includes you in six month's time).
Args:
type_name (str): The class of parameters to search for. Must be one of ``["Parameter", "SimulationParameter", "State"]`` or another future subclass of :py:class:`.ParameterBase`
family (Union[str, Tuple[str]]): A string or list or tuple of strings, that define one or more attribute families to search for
Returns:
dict: A nested dictionary of attributes that match the provided `type_name` and `family`
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
# - Filter own attribute dictionary by type key
matching_attributes = {
k: v for (k, v) in __registered_attributes.items() if v[1] == type_name
}
# - Filter by family
if family is not None:
if not isinstance(family, (tuple, list)):
family = (family,)
list_attributes = [
{k: v for (k, v) in matching_attributes.items() if v[2] is f}
for f in family
]
matching_attributes = dict(ChainMap(*list_attributes))
# - Just take values using getattr
matching_attributes = {k: getattr(self, k) for k in matching_attributes.keys()}
# - Append sub-module attributes as nested dictionaries
submodule_attributes = {}
for k, m in __modules.items():
mod_attributes = m[0]._get_attribute_family(type_name, family)
if (family and mod_attributes) or (not family):
submodule_attributes[k] = mod_attributes
# - Push submodule attributes into dictionary
if family and submodule_attributes or not family:
matching_attributes.update(submodule_attributes)
# - Return nested attributes
return matching_attributes
def attributes_named(self, name: Union[Tuple[str], List[str], str]) -> dict:
"""
Search for attributes of this or submodules by time
Args:
name (Union[str, Tuple[str]): The name of the attribute to search for
Returns:
dict: A nested dictionary of attributes that match `name`
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
# - Check if we were given a tuple or not
if not isinstance(name, (tuple, list)):
name = (name,)
# - Filter own attribute dictionary by name keys
list_attributes = [
{k: v for (k, v) in __registered_attributes.items() if k == n} for n in name
]
matching_attributes = dict(ChainMap(*list_attributes))
# - Just take values
matching_attributes = {k: v[0] for (k, v) in matching_attributes.items()}
# - Append sub-module attributes as nested dictionaries
submodule_attributes = {}
for k, m in __modules.items():
mod_attributes = m[0].attributes_named(name)
if mod_attributes:
submodule_attributes[k] = mod_attributes
# - Push submodule attributes into dictionary
if submodule_attributes:
matching_attributes.update(submodule_attributes)
# - Return nested attributes
return matching_attributes
def parameters(self, family: Union[str, Tuple, List] = None) -> Dict:
"""
Return a nested dictionary of module and submodule Parameters
Use this method to inspect the Parameters from this and all submodules. The optional argument `family` allows you to search for Parameters in a particular family — for example ``"weights"`` for all weights of this module and nested submodules.
Although the `family` argument is an arbitrary string, reasonable choises are ``"weights"``, ``"taus"`` for time constants, ``"biases"`` for biases...
Examples:
Obtain a dictionary of all Parameters for this module (including submodules):
>>> mod.parameters()
dict{ ... }
Obtain a dictionary of Parameters from a particular family:
>>> mod.parameters("weights")
dict{ ... }
Args:
family (str): The family of Parameters to search for. Default: ``None``; return all parameters.
Returns:
dict: A nested dictionary of Parameters of this module and all submodules
"""
return self._get_attribute_family("Parameter", family)
def simulation_parameters(self, family: Union[str, Tuple, List] = None) -> Dict:
"""
Return a nested dictionary of module and submodule SimulationParameters
Use this method to inspect the SimulationParameters from this and all submodules. The optional argument `family` allows you to search for SimulationParameters in a particular family.
Examples:
Obtain a dictionary of all SimulationParameters for this module (including submodules):
>>> mod.simulation_parameters()
dict{ ... }
Args:
family (str): The family of SimulationParameters to search for. Default: ``None``; return all SimulationParameter attributes.
Returns:
dict: A nested dictionary of SimulationParameters of this module and all submodules
"""
return self._get_attribute_family("SimulationParameter", family)
def state(self, family: Union[str, Tuple, List] = None) -> Dict:
"""
Return a nested dictionary of module and submodule States
Use this method to inspect the States from this and all submodules. The optional argument `family` allows you to search for States in a particular family.
Examples:
Obtain a dictionary of all States for this module (including submodules):
>>> mod.state()
dict{ ... }
Args:
family (str): The family of States to search for. Default: ``None``; return all State attributes.
Returns:
dict: A nested dictionary of States of this module and all submodules
"""
return self._get_attribute_family("State", family)
def modules(self) -> Dict:
"""
Return a dictionary of all sub-modules of this module
Returns:
dict: A dictionary containing all sub-modules. Each item will be named with the sub-module name.
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
return collections.OrderedDict([(k, m[0]) for (k, m) in __modules.items()])
def _reset_attribute(self, name: str) -> "ModuleBase":
"""
Reset an attribute to its initialisation value
Args:
name (str): The name of the attribute to reset
Returns:
self (`Module`): For compatibility with the functional API
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
# - Check that the attribute is registered
if name not in __registered_attributes:
raise KeyError(f"{name} is not a registered attribute.")
# - Get the initialisation function from the registry
(_, _, family, init_func, shape) = __registered_attributes[name]
# - Use the registered initialisation function, if present
if init_func is not None:
setattr(self, name, init_func(shape))
return self
def _has_registered_attribute(self, name: str) -> bool:
"""
Check if the module has a registered attribute
Args:
name (str): The name of the attribute to check
Returns:
bool: ``True`` if the attribute `name` is in the attribute registry, ``False`` otherwise.
"""
__registered_attributes, _ = self._get_attribute_registry()
return name in __registered_attributes
def reset_state(self) -> "ModuleBase":
"""
Reset the state of this module
Returns:
Module: The updated module is returned for compatibility with the functional API
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
# - Get a list of states
states = self.state()
# - Set self attributes
for k, v in __registered_attributes.items():
if k in states:
self._reset_attribute(k)
# - Reset submodule states
for k, m in __modules.items():
m[0] = m[0].reset_state()
return self
def reset_parameters(self):
"""
Reset all parameters in this module
Returns:
Module: The updated module is returned for compatibility with the functional API
"""
# - Get attribute registry
__registered_attributes, __modules = self._get_attribute_registry()
# - Get a list of parameters
parameters = self.parameters()
# - Set self attributes
for k, v in __registered_attributes.items():
if k in parameters:
self._reset_attribute(k)
# - Reset submodule states
for k, m in __modules.items():
m[0] = m[0].reset_parameters()
return self
@property
def class_name(self) -> str:
"""str: Class name of ``self``"""
# - Determine class name by removing "<class '" and "'>" and the package information
return type(self).__name__
@property
def name(self) -> str:
"""str: The name of this module, or an empty string if ``None``"""
try:
name = super().__getattribute__("_name")
return f"'{name}'" if name else ""
except:
return ""
@property
def full_name(self) -> str:
"""str: The full name of this module (class plus module name)"""
return f"{self.class_name} {self.name}"
@property
def spiking_input(self) -> bool:
"""bool: If ``True``, this module receives spiking input. If ``False``, this module expects continuous input."""
return self._spiking_input
@property
def spiking_output(self):
"""bool: If ``True``, this module sends spiking output. If ``False``, this module sends continuous output."""
return self._spiking_output
@property
def shape(self) -> tuple:
"""tuple: The shape of this module"""
return self._shape
@property
def size(self) -> int:
"""int: (DEPRECATED) The output size of this module"""
warn(
"The `size` property is deprecated. Please use `size_out` instead.",
DeprecationWarning,
)
return self._shape[-1]
@property
def size_out(self) -> int:
"""int: The output size of this module"""
return self._shape[-1]
@property
def size_in(self) -> int:
"""int: The input size of this module"""
return self._shape[0]
@abstractmethod
def evolve(self, input_data, record: bool = False) -> Tuple[Any, Any, Any]:
"""
Evolve the state of this module over input data
NOTE: THIS MODULE CLASS DOES NOT PROVIDE DOCUMENTATION FOR ITS EVOLVE METHOD. PLEASE UPDATE THE DOCUMENTATION FOR THIS MODULE.
Args:
input_data: The input data with shape ``(T, size_in)`` to evolve with
record (bool): If ``True``, the module should record internal state during evolution and return the record. If ``False``, no recording is required. Default: ``False``.
Returns:
tuple: (output, new_state, record)
output (np.ndarray): The output response of this module with shape ``(T, size_out)``
new_state (dict): A dictionary containing the updated state of this and all submodules after evolution
record (dict): A dictionary containing recorded state of this and all submodules, if requested using the `record` argument
"""
return None, None, None
def __call__(self, input_data, *args, **kwargs):
"""
Evolve the state of this module over input data
Args:
input_data: The input data with shape ``(T, size_in)`` to evolve this module with
*args: Additional positional arguments
**kwargs: Additional keyword arguments
Returns:
"""
# - Catch the case where we have been called with the raw output of a previous call
if isinstance(input_data, tuple) and len(input_data) == 3:
input_data, new_state, recorded_state = input_data
outputs, this_new_state, this_recorded_state = self.evolve(
input_data, *args, **kwargs
)
new_state.update({self.name: this_new_state})
recorded_state.update({self.name: this_recorded_state})
else:
outputs, new_state, recorded_state = self.evolve(
input_data, *args, **kwargs
)
return outputs, new_state, recorded_state
def _wrap_recorded_state(
self, recorded_dict: dict, t_start: float
) -> Dict[str, TimeSeries]:
"""
Convert a recorded dictionary to a :py:class:`TimeSeries` representation
This method is optional, and is provided to make the :py:meth:`timed` conversion to a :py:class:`TimedModule` work better. You should override this method in your custom :py:class:`Module`, to wrap each element of your recorded state dictionary as a :py:class:`TimeSeries`
Args:
state_dict (dict): A recorded state dictionary as returned by :py:meth:`.evolve`
t_start (float): The initial time of the recorded state, to use as the starting point of the time series
Returns:
Dict[str, TimeSeries]: The mapped recorded state dictionary, wrapped as :py:class:`TimeSeries` objects
"""
return recorded_dict
def as_graph(self) -> GraphModuleBase:
"""
Convert this module to a computational graph
Returns:
GraphModuleBase: The computational graph corresponding to this module
Raises:
NotImplementedError: If :py:meth:`.as_graph` is not implemented for this subclass
"""
raise NotImplementedError(
f"The Module class '{type(self).__name__}' used by module [{self}] does not implement the graph serialisation API"
)
class PostInitMetaMixin(type(ModuleBase)):
"""
A mixin base class that adds a ``__post_init__()`` method to a class. ``__post_init__()`` is called after the ``__init__()`` method, on instantiation of an object.
"""
def __call__(cls, *args, **kwargs):
# - Instantiate the object
obj = super().__call__(*args, **kwargs)
# - Check for a `__post_init__` method
if hasattr(cls, "__post_init__"):
cls.__post_init__(obj)
# - Clear the init flag
obj._in_Module_init = False
return obj
[docs]
class Module(ModuleBase, ABC, metaclass=PostInitMetaMixin):
"""
The native Python / numpy :py:class:`.Module` base class for Rockpool
This class acts as the base class for all "native" modules in Rockpool. To get started with using and writing your own Rockpool modules, see :ref:`/basics/getting_started.ipynb` and :ref:`/in-depth/api-low-level.ipynb`.
If you plan to write modules using Jax or Torch backends, you should use either :py:class:`.JaxModule` or :py:class:`.TorchModule` as base classes, respectively.
To get started with the Jax backend, see :ref:`/in-depth/api-functional.ipynb` and :ref:`/in-depth/jax-training.ipynb`.
To get started with the Torch backend, see :ref:`/in-depth/torch-api.ipynb` and :ref:`/in-depth/torch-training.ipynb`.
"""
[docs]
def timed(self, output_num: int = 0, dt: float = None, add_events: bool = False):
"""
Convert this module to a :py:class:`.TimedModule`
Args:
output_num (int): Specify which output of the module to take, if the module returns multiple output series. Default: ``0``, take the first (or only) output.
dt (float): Used to provide a time-step for this module, if the module does not already have one. If ``self`` already defines a time-step, then ``self.dt`` will be used. Default: ``None``
add_events (bool): Iff ``True``, the :py:class:`.TimedModule` will add events occurring on a single timestep on input and output. Default: ``False``, don't add time steps.
Returns: :py:class:`.TimedModule`: A timed module that wraps this module
"""
from rockpool.nn.modules.timed_module import TimedModuleWrapper
return TimedModuleWrapper(
self, output_num=output_num, dt=dt, add_events=add_events
)
[docs]
def _auto_batch(
self,
data: np.ndarray,
states: Tuple = (),
target_shapes: Tuple = None,
) -> Tuple[np.ndarray, Tuple[np.ndarray]]:
"""
Automatically replicate states over batches and verify input dimensions
Examples:
>>> data, (state0, state1, state2) = self._auto_batch(data, (self.state0, self.state1, self.state2))
This will verify that ``data`` has the correct final dimension (i.e. ``self.size_in``).
If ``data`` has only two dimensions ``(T, Nin)``, then it will be augmented to ``(1, T, Nin)``. The individual states will be replicated out from shape ``(a, b, c, ...)`` to ``(n_batches, a, b, c, ...)`` and returned.
If ``data`` has only a single dimension ``(T,)``, it will be expanded to ``(1, T, self.size_in)``.
``state0``, ``state1``, ``state2`` will be replicated out along the batch dimension.
>>> data, (state0,) = self._auto_batch(data, (self.state0,), ((10, -1, self.size_in),))
Attempt to replicate ``state0`` to a specified size ``(10, -1, self.size_in)``.
Args:
data (np.ndarray): Input data tensor. Either ``(batches, T, Nin)`` or ``(T, Nin)``
states (Tuple): Tuple of state variables. Each will be replicated out over batches by prepending a batch dimension
target_shapes (Tuple): A tuple of target size tuples, each corresponding to each state argument. The individual states will be replicated out to match the corresponding target sizes. If not provided (the default), then states will be only replicated along batches.
Returns:
(np.ndarray, Tuple[np.ndarray]) data, states
"""
# - Ensure data is a float tensor
data = np.array(data, "float")
# - Verify input data shape
if len(data.shape) == 1:
data = np.expand_dims(data, 0)
data = np.expand_dims(data, 2)
elif len(data.shape) == 2:
data = np.expand_dims(data, 0)
if data.shape[-1] == 1:
data = np.broadcast_to(data, (data.shape[0], data.shape[1], self.size_in))
# - Get shape of input
(n_batches, time_steps, n_connections) = data.shape
# - Check input dimensions
if n_connections != self.size_in:
raise ValueError(
"Input has wrong neuron dimension. It is {}, must be {}".format(
n_connections, self.size_in
)
)
# - Get target shapes
if target_shapes is None:
target_shapes = tuple(s.shape for s in states)
else:
target_shapes = tuple(
s.shape if shape is None else shape
for s, shape in zip(states, target_shapes)
)
# - Replicate shapes and return
states = tuple(
np.ones((n_batches, *shape)) * s for s, shape in zip(states, target_shapes)
)
return data, states