Unified OperationsΒΆ

BrainPy is targeted on multiple backends. Flexible switch between various backends needs to solve the problem of how to unify the different operations between them.

[1]:
import brainpy as bp

Intrinsically, BrainPy needs several necessary operations for numerical solvers, dynamics simulation and data construction.

[2]:
# necessary operations for numerical solvers

bp.ops.OPS_FOR_SOLVER
[2]:
['normal', 'sum', 'exp', 'shape']
[3]:
# necessary operations for neurodynamics simulation

bp.ops.OPS_FOR_SIMULATION
[3]:
['as_tensor', 'zeros', 'ones', 'arange', 'concatenate', 'where', 'reshape']
[4]:
# necessary data types

bp.ops.OPS_OF_DTYPE
[4]:
['bool', 'int', 'int32', 'int64', 'float', 'float32', 'float64']

However, if want to unify more commonly used operations, users can use brainpy.ops.set_buffer(backend, **operations) to set operation buffers.

For example, if users want to implement unified clip and square operations across different backends, ones can define like this:

[5]:
# NumPy

import numpy as np

bp.ops.set_buffer('numpy', clip=np.clip, sqrt=np.sqrt)
[6]:
# PyTorch

try:
    import torch

    bp.ops.set_buffer('pytorch', clip=torch.clamp, sqrt=torch.sqrt)

except ModuleNotFoundError:
    pass
[7]:
# TensorFlow

try:
    import tensorflow as tf

    bp.ops.set_buffer('tensorflow', clip=tf.clip_by_value, sqrt=tf.math.sqrt)

except ModuleNotFoundError:
    pass
[8]:
# Numba

try:
    import numba as nb

    @nb.njit
    def nb_clip(x, x_min, x_max):
        x = np.maximum(x, x_min)
        x = np.minimum(x, x_max)
        return x

    bp.ops.set_buffer('numba', clip=nb_clip, sqrt=np.sqrt)
    bp.ops.set_buffer('numba-parallel', clip=nb_clip, sqrt=np.sqrt)

except ModuleNotFoundError:
    pass
[9]:
# Numba-CUDA

try:
    import math
    import numba as nb
    from numba import cuda

    @cuda.jit(devicde=True)
    def cuda_clip(x, x_min, x_max):
        if x < x_min: return x_min
        elif x > x_max: return x_max
        else: return x

    bp.ops.set_buffer('numba-cuda', clip=nb_clip, sqrt=math.sqrt)

except ModuleNotFoundError:
    pass

After the buffer setting, users can use the unified operation to define models which will automatically works natively with buffered backends.

[10]:
def test(arr):
    return bp.ops.sqrt(bp.ops.clip(arr, 0., 1.))
[11]:
bp.backend.set('numpy')

test(bp.ops.as_tensor([-1, 0.5, 2.]))
[11]:
array([0.        , 0.70710678, 1.        ])
[12]:
bp.backend.set('pytorch')

test(bp.ops.as_tensor([-1, 0.5, 2.]))
[12]:
tensor([0.0000, 0.7071, 1.0000])
[13]:
bp.backend.set('tensorflow')

test(bp.ops.as_tensor([-1, 0.5, 2.]))
[13]:
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([0.        , 0.70710677, 1.        ], dtype=float32)>
[14]:
bp.backend.set('numba')

test(bp.ops.as_tensor([-1, 0.5, 2.]))
[14]:
array([0.        , 0.70710678, 1.        ])

Author: