This page was generated from docs/developer/backend-management.ipynb. Interactive online version:

Using the backend management system

Rockpool supports multiple optional computational backends, while keeping the number of core dependencies low. To do so we use a particular sub-package structure to isolate the backends, and provide some utility functions to assist with backend detection and management.

Code architecture of backend-specific Module packages

Rockpool keeps backend-specific Module subclasses in separate packages under nn.modules. This allows us to “quarantine” the backend to that particular sub-package.

Here we use the example of the nn.modules.torch sub-package.

Each new Module is implemented in its own Python file, which is written cleanly by assuming that all required dependencies are available. i.e., each file simply writes import torch without needing to perform any dependency checks.

This simplifies the code, but implies that if we attempt to import the symbols when a dependency is missing, we will raise an ImportError.

In __init__.py, we import all the Module classes up to the level of __init__.py to make the logical hierarchy of nn.modules simplier for the end-user. But this implies that we need to perform dependency checks in __init__.py, and handle andy missing dependencies accordingly.

We use the backend_available() function provided by utilities.backend_management to detect whether a given backend is importable. If it is missing, we create a fake class which raises an error if used, indiciating that the backend is not available and the class is not able to be used.

try:
    from .rate_torch import *
    ...

except:
    from rockpool.utilities.backend_management import (
        backend_available,
        missing_backend_shim,
    )
    if not backend_available('torch'):
        RateTorch = missing_backend_shim('RateTorch', 'torch')
        ...

    else:
        raise

This block attemps to import all exported symbols from rate_torch. If an error is raised, then backend_available() is used to check the availability of torch. If torch is available, this returns True.

If torch is missing, then a fake “shim” class RateTorch is created. The arguments to missing_backend_shim() specify the name of the shim class and the backend dependency or dependencies which are missing.

from rockpool.utilities.backend_management import missing_backend_shim

FakeClass = missing_backend_shim('FakeClass', 'torch')

FakeClass()

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
<ipython-input-1-d4a57251a9b4> in <module>
      3 FakeClass = missing_backend_shim('FakeClass', 'torch')
      4
----> 5 FakeClass()

~/SynSense Dropbox/Dylan Muir/LiveSync/Development/rockpool_GIT/rockpool/utilities/backend_management.py in __init__(self, *args, **kwargs)
    137         def __init__(self, *args, **kwargs):
    138             raise ModuleNotFoundError(
--> 139                 f"Missing the `{backend_name}` backend. `{class_name}` objects, and others relying on `{backend_name}` are not available."
    140             )
    141

ModuleNotFoundError: Missing the `torch` backend. `FakeClass` objects, and others relying on `torch` are not available.

Facilities provided by backend_management

Checking standard back-ends

“Standard” back-ends can be checked conveniently with backend_available(), which accepts the name of a back-end to check. “Standard” back-ends are {"numpy", "numba", "nest", "jax", "torch", "sinabs", "sinabs-exodus, "brian"}.

def backend_available(*backends) -> bool:
    ...

Checking non-standard back-ends

Other back-ends not included in the standard list can also easily be added and checked, without hacking the backend management package, by using the check_backend() function.

def check_backend(
    backend_name: str,
    required_modules: Optional[Union[Tuple[str], List[str]]] = None,
    check_flag: bool = True,
) -> bool:
    ...

Here backend_name is an arbitrary user-facing string specifying the back-end in a “nice” way. required_modules is a list of strings that will be each attempted to be imported. These specify the required python modules which comprise the back-end. If any of these modules cannot be imported for any reason, then check_backend() will return False. Otherwise it will return True.

check_flag provides a facility to perform a developer-defined arbitrary check for the backend, which is then incorporated in the back-end checking process. If check_flag = False then check_backend() will always return False.