This page was generated from docs/developer/UML-diagrams.ipynb. Interactive online version: Binder badge

UML diagrams for Rockpool

This notebook contains UML diagrams for the main packages and SW architecture of Rockpool.

[90]:

# - Load UML diagrams image
from PIL import Image
from pathlib import Path
im = Image.open(Path('uml', 'UML-diagrams.png'))

Module base and core derived classes, and combinators

The Module base class as well as all derived classes are under the nn.modules package. All user-defined modules as well as back-end base classes derive from Module or ModuleBase. ModuleBase is not designed as a user-facing class, and defines most of the features of Module. Module adds a facility for a post-initialisation callback on instantiation, used by some derived base classes.

Combinators used to build networks are in the package nn.combinators. These are adapted within nn.combinators for specific computational back-ends.

[93]:

im.crop([0, 0, 2450, 1800])

[93]:
../_images/developer_UML-diagrams_4_0.png

Jax and Torch core modules

Derived Modules for specific back-ends are under nn.modules, in sub-packages named by backend. Some backends require derived base classes to operate; e.g. JaxModule and TorchModule. All derived modules using jax inherit from JaxModule, and likewise for torch.

TorchModule also supports direct conversion of torch modules imported from torch.nn.Module.

[75]:

im.crop([2700, 0, 4750, 1400])

[75]:
../_images/developer_UML-diagrams_7_0.png

Rockpool Parameter classes

Classes for managing parameters are in the package nn.parameters. Parameter, State and SimulationParameter derive from ParameterBase.

Constant() is a special class which wraps any other class, permitting Rockpool to identify that a constant argument should not be trainable.

[84]:

im.crop([7450, 0, 9000, 800])

[84]:
../_images/developer_UML-diagrams_11_0.png

Hardware support packages

All hardware support packages are placed under devices. At present only devices.xylo is supported. See the documentation for Xylo for more information about these classes.

Different hardware versions are supported by providing sub-packages for each device. Within these subpackages are duplicate-named classes which support the corresponding specific device version.

e.g. devices.xylo.syns61201 provides XyloSim, XyloSamna, AFESamna etc. for the SYNS61201 device.

[87]:

im.crop([5000, 0, 7000, 1200])

[87]:
../_images/developer_UML-diagrams_14_0.png

Computational graph representation

All classes used to serialise and represent Rockpool modules in a computational graph are under the package graph.

graph contains the base class GraphModule, from which all derived GraphModule classes should inherit. It also provides GraphHolder and GraphNode classes.

The sub-package graph.utils provides a number of utility functions for manipulating graphs

[89]:

im.crop([0, 3500, 2300, 5750])

[89]:
../_images/developer_UML-diagrams_17_0.png

Training and transformations

The package training contains several sub-packages with useful functions for defining losses, regulariation, and training approaches.

The package transform contains sub-packages for parameter and network transformations. Currently only the package transform.quantize_methods is supported, providing methods for post-traning quantization for Xylo.

The beta implementation of QAT functionality for torch modules is under transform.torch_transform.

[79]:

im.crop([2400, 3500, 4600, 5200])

[79]:
../_images/developer_UML-diagrams_20_0.png