spyx.experimental#

Module Contents#

Classes#

PSU_LIF

Base class for Haiku modules.

StochasticAssociativeLIF

Base class for Haiku modules.

StochasticAssociativeCuBaLIF

Base class for Haiku modules.

SPSN

Prototype implementation of Stochastic Parallelizable Spiking Neuron:

Functions#

sigmoid_bernoulli([k, threshold, max_prob])

refractory_sigmoid_bernoulli([k, threshold])

spyx.experimental.sigmoid_bernoulli(k=10, threshold=1.0, max_prob=0.8)[source]#
spyx.experimental.refractory_sigmoid_bernoulli(k=50, threshold=1)[source]#
class spyx.experimental.PSU_LIF(hidden_shape, threshold=1, k=2, spike=True, name='PSULIF')[source]#

Bases: haiku.Module

Base class for Haiku modules.

A Haiku module is a lightweight container for variables and other modules. Modules typically define one or more “forward” methods (e.g. __call__) which apply operations combining user input and module parameters.

Modules must be initialized inside a transform() call.

For example:

>>> class AddModule(hk.Module):
...   def __call__(self, x):
...     w = hk.get_parameter("w", [], init=jnp.ones)
...     return x + w
>>> def forward_fn(x):
...   mod = AddModule()
...   return mod(x)
>>> forward = hk.transform(forward_fn)
>>> x = 1.
>>> rng = None
>>> params = forward.init(rng, x)
>>> print(forward.apply(params, None, x))
2.0
__call__(x)[source]#
class spyx.experimental.StochasticAssociativeLIF(hidden_shape, threshold=1, k=100, spike=True, name='SALIF')[source]#

Bases: haiku.Module

Base class for Haiku modules.

A Haiku module is a lightweight container for variables and other modules. Modules typically define one or more “forward” methods (e.g. __call__) which apply operations combining user input and module parameters.

Modules must be initialized inside a transform() call.

For example:

>>> class AddModule(hk.Module):
...   def __call__(self, x):
...     w = hk.get_parameter("w", [], init=jnp.ones)
...     return x + w
>>> def forward_fn(x):
...   mod = AddModule()
...   return mod(x)
>>> forward = hk.transform(forward_fn)
>>> x = 1.
>>> rng = None
>>> params = forward.init(rng, x)
>>> print(forward.apply(params, None, x))
2.0
__call__(key, x)[source]#
class spyx.experimental.StochasticAssociativeCuBaLIF(hidden_shape, threshold=1, k=100, name='SACuBaLIF')[source]#

Bases: haiku.Module

Base class for Haiku modules.

A Haiku module is a lightweight container for variables and other modules. Modules typically define one or more “forward” methods (e.g. __call__) which apply operations combining user input and module parameters.

Modules must be initialized inside a transform() call.

For example:

>>> class AddModule(hk.Module):
...   def __call__(self, x):
...     w = hk.get_parameter("w", [], init=jnp.ones)
...     return x + w
>>> def forward_fn(x):
...   mod = AddModule()
...   return mod(x)
>>> forward = hk.transform(forward_fn)
>>> x = 1.
>>> rng = None
>>> params = forward.init(rng, x)
>>> print(forward.apply(params, None, x))
2.0
__call__(key, u)[source]#
class spyx.experimental.SPSN(hidden_shape: tuple, threshold=1, k=10, name='SPSN')[source]#

Bases: haiku.Module

Prototype implementation of Stochastic Parallelizable Spiking Neuron:

https://doi.org/10.48550/arXiv.2306.12666

Currently only has Sigmoid-Bernoulli as a firing activation option. Needs further debugging.

Currently the output layer can only have a single beta value, and it tends to optimize to 0 which is odd.

Parameters:

hidden_shape (tuple) –

__call__(key, x)[source]#
X:

input tensor coming from previous layer. [Batch, Time, channels …]

Key:

JAX PRNGKey for computing stochastic spikes based on voltage potentials.