diff --git a/src/lava/proc/lif/ncmodels.py b/src/lava/proc/lif/ncmodels.py index 860a6572d..e68fe46b8 100644 --- 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,14 +119,14 @@ 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, @@ -134,8 +134,7 @@ def allocate(self, net: NetL2): ) # 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,7 +204,7 @@ 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, @@ -213,8 +212,7 @@ def allocate(self, net: NetL2): ) # 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)