Source code for norse.torch.module.iaf

import torch

from norse.torch.module.snn import SNNCell
from norse.torch.functional.iaf import (

[docs]class IAFCell(SNNCell):
[docs] def __init__(self, p: IAFParameters = IAFParameters(), dt: float = 0.001): r"""Feedforward step of an integrate-and-fire neuron, computing a single step .. math:: \dot{v} = v together with the jump condition .. math:: z = \Theta(v - v_{\text{th}}) and transition equation .. math:: v = (1-z) v + z v_{\text{reset}} Parameters: p (IAFParameters): parameters of a leaky integrate and fire neuron dt (float): Integration timestep to use (unused, but added for compatibility) """ super().__init__(iaf_feed_forward_step, self.initial_state, p, dt)
def initial_state(self, input_tensor: torch.Tensor) -> IAFFeedForwardState: state = IAFFeedForwardState( v=torch.full( input_tensor.shape, torch.as_tensor(self.p.v_reset.detach()), device=input_tensor.device, dtype=torch.float32, ), ) state.v.requires_grad = True return state