changed Node to Nodes (#195)

* changed Node to Nodes
lava-nc · Jul 4, 2022 · fbb9dc6 · fbb9dc6
1 parent 21c31ce
commit fbb9dc6
Showing 1 changed file with 10 additions and 12 deletions.
diff --git a/src/lava/proc/lif/ncmodels.py b/src/lava/proc/lif/ncmodels.py
@@ -15,7 +15,7 @@
     class NetL2:
         pass
 
-from lava.magma.compiler.node import Node
+from lava.magma.core.nets.tables import Nodes
 from lava.magma.core.resources import NeuroCore, Loihi2NeuroCore
 from lava.magma.core.decorator import implements, requires, tag
 from lava.magma.core.model.nc.model import AbstractNcProcessModel
@@ -48,7 +48,7 @@ def allocate(self, net: NetL2):
         flat_size = np.product(list(self.proc_params['shape']))
 
         # Allocate neurons
-        neurons_cfg: Node = net.neurons_cfg.allocate(
+        neurons_cfg: Nodes = net.neurons_cfg.allocate(
             shape=1,
             du=self.du,
             dv=self.dv,
@@ -64,7 +64,7 @@ def allocate(self, net: NetL2):
             stack_in=0,
             refractory_delay=1,
             bap_action=0)
-        neurons: Node = net.neurons.allocate_hcode(
+        neurons: Nodes = net.neurons.allocate_hcode(
             shape=flat_size,
             u=self.u,
             v=self.v,
@@ -73,7 +73,7 @@ def allocate(self, net: NetL2):
             vth=None)
 
         # Allocate output axons
-        ax_out: Node = net.ax_out.allocate(
+        ax_out: Nodes = net.ax_out.allocate(
             shape=flat_size, num_message_bits=0)
 
         # Connect InPort of Process to neurons
@@ -119,23 +119,22 @@ def allocate(self, net: NetL2):
         vth_reg = np.left_shift(self.vth.var.init, 6)
 
         # Allocate neurons
-        neurons_cfg: Node = net.neurons_cfg.allocate(
+        neurons_cfg: Nodes = net.neurons_cfg.allocate(
             shape=(1, ),
             ucode=ucode_file,
             vth=vth_reg,
             du=4096 - self.du.var.get(),
             dv=4096 - self.dv.var.get()
         )
-        neurons: Node = net.neurons.allocate_ucode(
+        neurons: Nodes = net.neurons.allocate_ucode(
             shape=shape,
             u=self.u,
             v=self.v,
             bias=self.bias_mant.var.get() * 2 ** (self.bias_exp.var.get()),
         )
 
         # Allocate output axons
-        ax_out: Node = net.ax_out.allocate(shape=shape,
-                                           num_message_bits=0)
+        ax_out: Nodes = net.ax_out.allocate(shape=shape, num_message_bits=0)
 
         # Connect InPort of Process to neurons
         self.a_in.connect(neurons)
@@ -196,7 +195,7 @@ def allocate(self, net: NetL2):
         interval_sub_1_reg = reset_interval - 1
 
         # Allocate neurons
-        neurons_cfg: Node = net.neurons_cfg.allocate(
+        neurons_cfg: Nodes = net.neurons_cfg.allocate(
             shape=(1, ),
             ucode=ucode_file,
             vth=vth_reg,
@@ -205,16 +204,15 @@ def allocate(self, net: NetL2):
             interval_sub_1=interval_sub_1_reg,
             offset=reset_offset
         )
-        neurons: Node = net.neurons.allocate_ucode(
+        neurons: Nodes = net.neurons.allocate_ucode(
             shape=shape,
             u=self.u,
             v=self.v,
             bias=self.bias_mant.var.get() * 2 ** (self.bias_exp.var.get()),
         )
 
         # Allocate output axons
-        ax_out: Node = net.ax_out.allocate(shape=shape,
-                                           num_message_bits=0)
+        ax_out: Nodes = net.ax_out.allocate(shape=shape, num_message_bits=0)
 
         # Connect InPort of Process to neurons
         self.a_in.connect(neurons)