added atan, hside, lif structure. lif update broken, needs fix.

src/lib.rs

@@ -1,4 +1,4 @@
-#[allow(non_snake_case)]
+#![allow(non_snake_case)]
 pub mod encoders;
 pub mod neurons;
 pub mod surrogate;

src/neurons/heaviside.rs

@@ -0,0 +1,13 @@
+#![allow(non_snake_case)]
+use crate::surrogate::SurrogateFn;
+use burn::prelude::*;
+
+pub fn heaviside<B: Backend, const D: usize>(
+    input: Tensor<B, D>,
+    thres: f32,
+    s: SurrogateFn,
+) -> Tensor<B, D> {
+    let y = input.clone().greater_equal_elem(thres).float();
+    let yb = s.forward(input);
+    (y - yb.clone().detach()).detach() + yb
+}

src/neurons/leaky.rs

@@ -1,24 +1,64 @@
-use burn::{Tensor, module::Module, prelude::Backend, tensor::backend::AutodiffBackend};
+#![allow(non_snake_case)]
 
+use crate::neurons::heaviside::heaviside;
+use crate::surrogate::SurrogateFn;
+use burn::module::Param;
+use burn::prelude::*;
+use burn::{Tensor, module::Module, prelude::Backend};
+
+#[derive(Config, Debug)]
+pub struct LIFConfig {
+    beta: f32,
+    threshold: f32,
+    neurons: usize,
+    spikeGrad: SurrogateFn,
+}
+
+impl LIFConfig {
+    pub fn init<B: Backend, const D: usize>(&self, device: &B::Device) -> LIF<B, D> {
+        let initMem = Param::from_tensor(Tensor::<B, D>::zeros([1, self.neurons], device));
+        LIF {
+            beta: self.beta,
+            threshold: self.threshold,
+            neurons: self.neurons,
+            hidden: initMem,
+        }
+    }
+}
+
+// TODO: tensor cloning and its lifecycle is probably wrong, may cause comp graph to drop. Refer burn example to find proper tensor handling..
 #[derive(Debug, Module)]
-pub struct LIF<B: Backend> {
-    pub v: Tensor<B, 1>,
+pub struct LIF<B: Backend, const D: usize> {
+    beta: f32,
+    threshold: f32,
+    neurons: usize,
+    pub hidden: Param<Tensor<B, D>>,
 }
 
-impl<B: Backend> LIF<B> {
-    pub fn new(device: &B::Device) -> Self {
-        Self {
-            v: Tensor::zeros([1], device),
-        }
-    }
-}
+impl<B: Backend, const D: usize> LIF<B, D> {
+    pub fn forward(&mut self, input: Tensor<B, D>) -> Tensor<B, D> {
+        // leaky and Integrate
+        //
+        let curMem = self.hidden.val();
+        let nxtMem = curMem.mul_scalar(self.beta).add(input.clone());
+        // fire
+        let spikes = heaviside::<B, D>(nxtMem, self.threshold, SurrogateFn::FastSigmoid);
 
-impl<B: AutodiffBackend> LIF<B> {
-    pub fn forward(&mut self, input: Tensor<B, 2>) -> Tensor<B, 2> {
-        input
+        self.hidden
+            .val()
+            .sub(spikes.clone().mul_scalar(self.threshold)); // requires update step fix. currently doesnt update.
+        spikes
     }
 
-    pub fn backward(&mut self, x: Tensor<B, 2>) -> Tensor<B, 2> {
-        x
+    fn mem_reset(&self, mem: Tensor<B, D>) -> Tensor<B, D> {
+        mem.sub_scalar(self.threshold)
+    }
+
+    pub fn reset(&mut self) {
+        // self.hidden = self.hidden.zeros_like();
+    }
+
+    pub fn init(&mut self, batch: usize, device: &B::Device) {
+        // self.hidden = Tensor::zeros(Shape::new([batch, self.neurons]), device);
     }
 }

src/neurons/mod.rs

@@ -1 +1,2 @@
+mod heaviside;
 mod leaky;

src/surrogate/atan.rs

@@ -0,0 +1,5 @@
+use burn::prelude::*;
+
+pub fn atan<B: Backend, const D: usize>(input: Tensor<B, D>) -> Tensor<B, D> {
+    (1.0 / std::f32::consts::PI) * (input.mul_scalar(std::f32::consts::PI).atan())
+}

src/surrogate/fast_sigmoid.rs

@@ -1 +1,5 @@
+use burn::prelude::*;
 
+pub fn fast_sigmoid<B: Backend, const D: usize>(input: Tensor<B, D>) -> Tensor<B, D> {
+    input.clone() / (1 + input.abs())
+}

src/surrogate/mod.rs

@@ -1 +1,21 @@
+mod atan;
 mod fast_sigmoid;
+
+use atan::atan;
+use burn::{Tensor, config::Config, tensor::backend::Backend};
+use fast_sigmoid::fast_sigmoid;
+
+#[derive(Config, Debug)]
+pub enum SurrogateFn {
+    FastSigmoid,
+    ATan,
+}
+
+impl SurrogateFn {
+    pub fn forward<B: Backend, const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
+        match self {
+            SurrogateFn::FastSigmoid => fast_sigmoid(input),
+            SurrogateFn::ATan => atan(input),
+        }
+    }
+}