tether.nn package

Submodules

tether.nn.alif module

class tether.nn.alif.ALIF(n_neurons, decay_v=0.9, decay_a=0.9, threshold=1.0, beta=0.5, alpha=2.0, store_traces=False)

Bases: Module

forward(x_seq)

Forward pass of the ALIF module.

Parameters:

x_seq (torch.Tensor) – Input sequence.

Returns:

Output spikes with the same shape as input.

Return type:

torch.Tensor

tether.nn.attention module

class tether.nn.attention.SpikingSelfAttention(dim, num_heads=8, decay=0.9, threshold=1.0)

Bases: Module

forward(x_seq)

Forward pass of the SpikingSelfAttention.

Parameters:

x_seq (torch.Tensor) – Input sequence of shape (T, B, N, D).

Returns:

Output sequence of shape (T, B, N, D).

Return type:

torch.Tensor

reset_states()

Reset the KV state buffer to zeros.

tether.nn.block module

class tether.nn.block.SpikingTransformerBlock(dim, num_heads=8, mlp_ratio=4)

Bases: Module

forward(x)

Forward pass of the SpikingTransformerBlock.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

torch.Tensor

tether.nn.lif module

class tether.nn.lif.LIF(n_neurons, decay=0.9, threshold=1.0, alpha=2.0, surrogate=None, store_traces=False)

Bases: Module

property alpha

forward(x_seq)

Forward pass of the LIF module.

Parameters:

x_seq (torch.Tensor) – Input sequence.

Returns:

Output spikes with the same shape as input.

Return type:

torch.Tensor

tether.nn.plif module

class tether.nn.plif.PLIF(n_neurons, init_decay=0.9, init_threshold=1.0, alpha=2.0, surrogate=None, store_traces=False)

Bases: Module

forward(x_seq)

Forward pass of the PLIF module.

Parameters:

x_seq (torch.Tensor) – Input sequence of shape (Time, Batch, Neurons) or (Time, Batch * Neurons).

Returns:

Output spikes with the same shape as input.

Return type:

torch.Tensor

tether.nn.surrogates module

class tether.nn.surrogates.Arctan(alpha=2.0, trainable=False)

Bases: Surrogate

Arctan surrogate gradient.

The surrogate derivative is given by:

\[\begin{split}f'(x) = \\frac{1}{1 + (\\alpha \\pi x)^2}\end{split}\]

where x is the normalized membrane potential (v - threshold).

get_id()

Return the integer ID corresponding to the surrogate type used in Triton kernels.

Returns:

Surrogate type ID.

Return type:

int

class tether.nn.surrogates.FastSigmoid(alpha=2.0, trainable=False)

Bases: Surrogate

Fast Sigmoid (approximated) surrogate gradient.

Uses a computationally cheaper approximation of the sigmoid derivative:

\[\begin{split}f'(x) = \\frac{1}{(1 + |\\alpha x|)^2}\end{split}\]

This avoids expensive exponential operations.

get_id()

Return the integer ID corresponding to the surrogate type used in Triton kernels.

Returns:

Surrogate type ID.

Return type:

int

class tether.nn.surrogates.Sigmoid(alpha=2.0, trainable=False)

Bases: Surrogate

Sigmoid surrogate gradient.

The surrogate function is a sigmoid, and its derivative is:

\[\begin{split}f'(x) = \\alpha \\cdot \\sigma(\\alpha x) \\cdot (1 - \\sigma(\\alpha x))\end{split}\]

where \(\\sigma\) is the logistic sigmoid function and x is the membrane potential gap.

get_id()

Return the integer ID corresponding to the surrogate type used in Triton kernels.

Returns:

Surrogate type ID.

Return type:

int

class tether.nn.surrogates.Surrogate(alpha=2.0, trainable=False)

Bases: Module

Base class for surrogate gradient functions used in Spiking Neural Networks.

Surrogate gradients allow for backpropagation through the non-differentiable Heaviside step function used for spike generation.

Parameters:

• alpha (float, optional) – Scaling parameter that controls the steepness/width of the surrogate derivative. Default is 2.0.

• trainable (bool, optional) – If True, alpha becomes a learnable parameter. Default is False.

get_id()

Return the integer ID corresponding to the surrogate type used in Triton kernels.

Returns:

Surrogate type ID.

Return type:

int

Module contents

class tether.nn.ALIF(n_neurons, decay_v=0.9, decay_a=0.9, threshold=1.0, beta=0.5, alpha=2.0, store_traces=False)

Bases: Module

forward(x_seq)

Forward pass of the ALIF module.

Parameters:

x_seq (torch.Tensor) – Input sequence.

Returns:

Output spikes with the same shape as input.

Return type:

torch.Tensor

class tether.nn.Arctan(alpha=2.0, trainable=False)

Bases: Surrogate

Arctan surrogate gradient.

The surrogate derivative is given by:

\[\begin{split}f'(x) = \\frac{1}{1 + (\\alpha \\pi x)^2}\end{split}\]

where x is the normalized membrane potential (v - threshold).

get_id()

Return the integer ID corresponding to the surrogate type used in Triton kernels.

Returns:

Surrogate type ID.

Return type:

int

class tether.nn.FastSigmoid(alpha=2.0, trainable=False)

Bases: Surrogate

Fast Sigmoid (approximated) surrogate gradient.

Uses a computationally cheaper approximation of the sigmoid derivative:

\[\begin{split}f'(x) = \\frac{1}{(1 + |\\alpha x|)^2}\end{split}\]

This avoids expensive exponential operations.

get_id()

Return the integer ID corresponding to the surrogate type used in Triton kernels.

Returns:

Surrogate type ID.

Return type:

int

class tether.nn.LIF(n_neurons, decay=0.9, threshold=1.0, alpha=2.0, surrogate=None, store_traces=False)

Bases: Module

property alpha

forward(x_seq)

Forward pass of the LIF module.

Parameters:

x_seq (torch.Tensor) – Input sequence.

Returns:

Output spikes with the same shape as input.

Return type:

torch.Tensor

class tether.nn.PLIF(n_neurons, init_decay=0.9, init_threshold=1.0, alpha=2.0, surrogate=None, store_traces=False)

Bases: Module

forward(x_seq)

Forward pass of the PLIF module.

Parameters:

x_seq (torch.Tensor) – Input sequence of shape (Time, Batch, Neurons) or (Time, Batch * Neurons).

Returns:

Output spikes with the same shape as input.

Return type:

torch.Tensor

class tether.nn.Sigmoid(alpha=2.0, trainable=False)

Bases: Surrogate

Sigmoid surrogate gradient.

The surrogate function is a sigmoid, and its derivative is:

\[\begin{split}f'(x) = \\alpha \\cdot \\sigma(\\alpha x) \\cdot (1 - \\sigma(\\alpha x))\end{split}\]

where \(\\sigma\) is the logistic sigmoid function and x is the membrane potential gap.

get_id()

Return the integer ID corresponding to the surrogate type used in Triton kernels.

Returns:

Surrogate type ID.

Return type:

int

class tether.nn.SpikingSelfAttention(dim, num_heads=8, decay=0.9, threshold=1.0)

Bases: Module

forward(x_seq)

Forward pass of the SpikingSelfAttention.

Parameters:

x_seq (torch.Tensor) – Input sequence of shape (T, B, N, D).

Returns:

Output sequence of shape (T, B, N, D).

Return type:

torch.Tensor

reset_states()

Reset the KV state buffer to zeros.

class tether.nn.SpikingTransformerBlock(dim, num_heads=8, mlp_ratio=4)

Bases: Module

forward(x)

Forward pass of the SpikingTransformerBlock.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

torch.Tensor

class tether.nn.Surrogate(alpha=2.0, trainable=False)

Bases: Module

Base class for surrogate gradient functions used in Spiking Neural Networks.

Surrogate gradients allow for backpropagation through the non-differentiable Heaviside step function used for spike generation.

Parameters:

• alpha (float, optional) – Scaling parameter that controls the steepness/width of the surrogate derivative. Default is 2.0.

• trainable (bool, optional) – If True, alpha becomes a learnable parameter. Default is False.

get_id()

Return the integer ID corresponding to the surrogate type used in Triton kernels.

Returns:

Surrogate type ID.

Return type:

int