refactored and passed test for deconv/stdp-deconv-syn and other minor cleanup for conv/deconv support

Alexander Ororbia · Alexander Ororbia · commit 6cb319a94289 · 2025-11-07T14:17:51.000-05:00
diff --git a/ngclearn/components/synapses/convolution/__init__.py b/ngclearn/components/synapses/convolution/__init__.py
@@ -2,7 +2,7 @@
 from .staticConvSynapse import StaticConvSynapse
 from .deconvSynapse import DeconvSynapse
 from .staticDeconvSynapse import StaticDeconvSynapse
-from .hebbianConvSynapse import HebbianConvSynapse
-from .hebbianDeconvSynapse import HebbianDeconvSynapse
+#from .hebbianConvSynapse import HebbianConvSynapse
+# from .hebbianDeconvSynapse import HebbianDeconvSynapse
 from .traceSTDPConvSynapse import TraceSTDPConvSynapse
 from .traceSTDPDeconvSynapse import TraceSTDPDeconvSynapse
diff --git a/ngclearn/components/synapses/convolution/convSynapse.py b/ngclearn/components/synapses/convolution/convSynapse.py
@@ -99,17 +99,15 @@ def __init__(
             dist.initialize_params(subkeys[2], bias_init, (1, shape[1])) if bias_init else 0.0
         )
 
-    # @transition(output_compartments=["outputs"])
-    # @staticmethod
     @compilable
     def advance_state(self): #Rscale, padding, stride, weights, biases, inputs):
         _x = self.inputs.get()
         ## FIXME: does resist_scale affect update rules?
-        outputs = conv2d(_x, self.weights.get(), stride_size=self.stride, padding=self.padding) * self.resist_scale + self.biases.get()
+        outputs = conv2d(
+            _x, self.weights.get(), stride_size=self.stride, padding=self.padding
+        ) * self.resist_scale + self.biases.get()
         self.outputs.set(outputs)
 
-    # @transition(output_compartments=["inputs", "outputs"])
-    # @staticmethod
     @compilable
     def reset(self): #in_shape, out_shape):
         preVals = jnp.zeros(self.in_shape)
diff --git a/ngclearn/components/synapses/convolution/deconvSynapse.py b/ngclearn/components/synapses/convolution/deconvSynapse.py
@@ -1,15 +1,14 @@
 from jax import random, numpy as jnp, jit
-from ngclearn.components.jaxComponent import JaxComponent
-from ngcsimlib.compilers.process import transition
-from ngcsimlib.component import Component
 from ngcsimlib.compartment import Compartment
-
+from ngcsimlib.parser import compilable
 from ngclearn.utils.weight_distribution import initialize_params
 from ngcsimlib.logger import info
-from ngclearn.utils import tensorstats
 import ngclearn.utils.weight_distribution as dist
 from ngclearn.components.synapses.convolution.ngcconv import deconv2d
 
+from ngclearn.components.jaxComponent import JaxComponent
+
+
 class DeconvSynapse(JaxComponent): ## base-level deconvolutional cable
     """
     A base deconvolutional (transposed convolutional) synaptic cable.
@@ -61,7 +60,7 @@ def __init__(
         self.shape = shape ## shape of synaptic filter tensor
         x_size, x_size = x_shape
         self.x_size = x_size
-        self.Rscale = resist_scale ## post-transformation scale factor
+        self.resist_scale = resist_scale ## post-transformation scale factor
         self.padding = padding
         self.stride = stride
 
@@ -70,7 +69,7 @@ def __init__(
         self.pad_args = None
 
         ######################### set up compartments ##########################
-        tmp_key, *subkeys = random.split(self.key.value, 4)
+        tmp_key, *subkeys = random.split(self.key.get(), 4)
         weights = dist.initialize_params(subkeys[0], filter_init,
                                          shape)  ## filter tensor
         self.batch_size = batch_size # 1
@@ -89,36 +88,35 @@ def __init__(
                                                          (1, shape[1]))
                                   if bias_init else 0.0)
 
-    @transition(output_compartments=["outputs"])
-    @staticmethod
-    def advance_state(Rscale, padding, stride, weights, biases, inputs):
-        _x = inputs
-        out = deconv2d(_x, weights, stride_size=stride, padding=padding) * Rscale + biases
-        return out
-
-    @transition(output_compartments=["inputs", "outputs"])
-    @staticmethod
-    def reset(in_shape, out_shape):
-        preVals = jnp.zeros(in_shape)
-        postVals = jnp.zeros(out_shape)
-        inputs = preVals
-        outputs = postVals
-        return inputs, outputs
-
-    def save(self, directory, **kwargs):
-        file_name = directory + "/" + self.name + ".npz"
-        if self.bias_init != None:
-            jnp.savez(file_name, weights=self.weights.value,
-                      biases=self.biases.value)
-        else:
-            jnp.savez(file_name, weights=self.weights.value)
-
-    def load(self, directory, **kwargs):
-        file_name = directory + "/" + self.name + ".npz"
-        data = jnp.load(file_name)
-        self.weights.set(data['weights'])
-        if "biases" in data.keys():
-            self.biases.set(data['biases'])
+    @compilable
+    def advance_state(self):
+        _x = self.inputs.get()
+        out = deconv2d(
+            _x, self.weights.get(), stride_size=self.stride, padding=self.padding
+        ) * self.resist_scale + self.biases.get()
+        self.outputs.set(out)
+
+    @compilable
+    def reset(self): #in_shape, out_shape):
+        preVals = jnp.zeros(self.in_shape)
+        postVals = jnp.zeros(self.out_shape)
+        self.inputs.set(preVals)
+        self.outputs.set(postVals)
+
+    # def save(self, directory, **kwargs):
+    #     file_name = directory + "/" + self.name + ".npz"
+    #     if self.bias_init != None:
+    #         jnp.savez(file_name, weights=self.weights.get(),
+    #                   biases=self.biases.get())
+    #     else:
+    #         jnp.savez(file_name, weights=self.weights.get())
+    #
+    # def load(self, directory, **kwargs):
+    #     file_name = directory + "/" + self.name + ".npz"
+    #     data = jnp.load(file_name)
+    #     self.weights.set(data['weights'])
+    #     if "biases" in data.keys():
+    #         self.biases.set(data['biases'])
 
     @classmethod
     def help(cls): ## component help function
@@ -151,17 +149,3 @@ def help(cls): ## component help function
                 "dynamics": "outputs = [K @.T inputs] * R + b",
                 "hyperparameters": hyperparams}
         return info
-
-    def __repr__(self):
-        comps = [varname for varname in dir(self) if Compartment.is_compartment(getattr(self, varname))]
-        maxlen = max(len(c) for c in comps) + 5
-        lines = f"[{self.__class__.__name__}] PATH: {self.name}\n"
-        for c in comps:
-            stats = tensorstats(getattr(self, c).value)
-            if stats is not None:
-                line = [f"{k}: {v}" for k, v in stats.items()]
-                line = ", ".join(line)
-            else:
-                line = "None"
-            lines += f"  {f'({c})'.ljust(maxlen)}{line}\n"
-        return lines
diff --git a/ngclearn/components/synapses/convolution/traceSTDPConvSynapse.py b/ngclearn/components/synapses/convolution/traceSTDPConvSynapse.py
@@ -146,17 +146,9 @@ def _compute_update(self): #pretrace_target, Aplus, Aminus, stride, pad_args, de
         dWeights = (dW_ltp + dW_ltd)
         return dWeights
 
-    # @transition(output_compartments=["weights", "dWeights"])
-    # @staticmethod
     @compilable
     def evolve(self):
-        # pretrace_target, Aplus, Aminus, w_decay, w_bound, stride, pad_args, delta_shape, preSpike, preTrace,
-        # postSpike, postTrace, weights, eta
-
         dWeights = self._compute_update()
-        # dWeights = TraceSTDPConvSynapse._compute_update(
-        #     pretrace_target, Aplus, Aminus, stride, pad_args, delta_shape, preSpike, preTrace, postSpike, postTrace
-        # )
         if self.w_decay > 0.:  ## apply synaptic decay
             weights = self.weights.get() + dWeights * self.eta - self.weights.get() * self.w_decay  ## conduct decayed STDP-ascent
         else:
@@ -169,10 +161,8 @@ def evolve(self):
         self.weights.set(weights)
         self.dWeights.set(dWeights)
 
-    # @transition(output_compartments=["dInputs"])
-    # @staticmethod
     @compilable
-    def backtransmit(self): # x_size, shape, stride, padding, x_delta_shape, antiPad, postSpike, weights
+    def backtransmit(self):
         ## action-backpropagating routine
         ## calc dInputs - adjustment w.r.t. input signal
         k_size, k_size, n_in_chan, n_out_chan = self.shape
diff --git a/ngclearn/components/synapses/convolution/traceSTDPDeconvSynapse.py b/ngclearn/components/synapses/convolution/traceSTDPDeconvSynapse.py
@@ -1,17 +1,14 @@
 from jax import random, numpy as jnp, jit
-from ngcsimlib.compilers.process import transition
-from ngcsimlib.component import Component
 from ngcsimlib.compartment import Compartment
-
-from .deconvSynapse import DeconvSynapse
+from ngcsimlib.parser import compilable
 from ngclearn.utils.weight_distribution import initialize_params
-from ngcsimlib.logger import info
-from ngclearn.utils import tensorstats
 import ngclearn.utils.weight_distribution as dist
+
+from ngclearn.components.synapses.convolution.deconvSynapse import DeconvSynapse
+
 from ngclearn.components.synapses.convolution.ngcconv import (deconv2d, _calc_dX_deconv,
                                                               _calc_dK_deconv, calc_dX_deconv,
                                                               calc_dK_deconv)
-from ngclearn.utils.optim import get_opt_init_fn, get_opt_step_fn
 
 class TraceSTDPDeconvSynapse(DeconvSynapse): ## trace-based STDP deconvolutional cable
     """
@@ -92,7 +89,7 @@ def __init__(
 
         ######################### set up compartments ##########################
         ## Compartment setup and shape computation
-        self.dWeights = Compartment(self.weights.value * 0)
+        self.dWeights = Compartment(self.weights.get() * 0)
         self.dInputs = Compartment(jnp.zeros(self.in_shape))
         self.preSpike = Compartment(jnp.zeros(self.in_shape))
         self.preTrace = Compartment(jnp.zeros(self.in_shape))
@@ -108,76 +105,73 @@ def __init__(
     def _init(self, batch_size, x_size, shape, stride, padding, pad_args, weights):
         k_size, k_size, n_in_chan, n_out_chan = shape
         _x = jnp.zeros((batch_size, x_size, x_size, n_in_chan))
-        _d = deconv2d(_x, self.weights.value, stride_size=self.stride,
+        _d = deconv2d(_x, self.weights.get(), stride_size=self.stride,
                       padding=self.padding) * 0
         _dK = _calc_dK_deconv(_x, _d, stride_size=self.stride, out_size=k_size)
         ## get filter update correction
-        dx = _dK.shape[0] - self.weights.value.shape[0]
-        dy = _dK.shape[1] - self.weights.value.shape[1]
+        dx = _dK.shape[0] - self.weights.get().shape[0]
+        dy = _dK.shape[1] - self.weights.get().shape[1]
         self.delta_shape = (abs(dx), abs(dy))
 
         ## get input update correction
-        _dx = _calc_dX_deconv(self.weights.value, _d, stride_size=self.stride,
+        _dx = _calc_dX_deconv(self.weights.get(), _d, stride_size=self.stride,
                        padding=self.padding)
         dx = (_dx.shape[1] - _x.shape[1])  # abs()
         dy = (_dx.shape[2] - _x.shape[2])
         self.x_delta_shape = (dx, dy)
 
-    @staticmethod
-    def _compute_update(
-            pretrace_target, Aplus, Aminus, shape, stride, padding, delta_shape, preSpike, preTrace, postSpike, postTrace
-    ):
-        k_size, k_size, n_in_chan, n_out_chan = shape
+    def _compute_update(self):
+        k_size, k_size, n_in_chan, n_out_chan = self.shape
         ## calc dFilters
-        dW_ltp = calc_dK_deconv(preTrace - pretrace_target, postSpike * Aplus,
-                                delta_shape=delta_shape, stride_size=stride,
-                                out_size=k_size, padding=padding)
-        dW_ltd = -calc_dK_deconv(preSpike, postTrace * Aminus,
-                                 delta_shape=delta_shape, stride_size=stride,
-                                 out_size=k_size, padding=padding)
+        dW_ltp = calc_dK_deconv(
+            self.preTrace.get() - self.pretrace_target, self.postSpike.get() * self.Aplus,
+            delta_shape=self.delta_shape, stride_size=self.stride, out_size=k_size, padding=self.padding
+        )
+        dW_ltd = -calc_dK_deconv(
+            self.preSpike.get(), self.postTrace.get() * self.Aminus, delta_shape=self.delta_shape,
+            stride_size=self.stride, out_size=k_size, padding=self.padding
+        )
         dWeights = (dW_ltp + dW_ltd)
         return dWeights
 
-    @transition(output_compartments=["weights", "dWeights"])
-    @staticmethod
-    def evolve(
-            pretrace_target, Aplus, Aminus, w_decay, w_bound, shape, stride, padding, delta_shape, preSpike, preTrace,
-            postSpike, postTrace, weights, eta
-    ):
-        dWeights = TraceSTDPDeconvSynapse._compute_update(
-            pretrace_target, Aplus, Aminus, shape, stride, padding, delta_shape,
-            preSpike, preTrace, postSpike, postTrace
-        )
-        if w_decay > 0.: ## apply synaptic decay
-            weights = weights + dWeights * eta - weights * w_decay  ## conduct decayed STDP-ascent
+    @compilable
+    def evolve(self):
+        dWeights = self._compute_update()
+        # dWeights = TraceSTDPDeconvSynapse._compute_update(
+        #     pretrace_target, Aplus, Aminus, shape, stride, padding, delta_shape,
+        #     preSpike, preTrace, postSpike, postTrace
+        # )
+        if self.w_decay > 0.: ## apply synaptic decay and conduct decayed STDP-ascent
+            weights = self.weights.get() + dWeights * self.eta - self.weights.get() * self.w_decay
         else:
-            weights = weights + dWeights * eta ## conduct STDP-ascent
+            weights = self.weights.get() + dWeights * self.eta ## conduct STDP-ascent
         ## Apply any enforced filter constraints
-        if w_bound > 0.: ## enforce non-negativity
+        if self.w_bound > 0.: ## enforce non-negativity
             eps = 0.01  # 0.001
-            weights = jnp.clip(weights, eps, w_bound - eps)
-        return weights, dWeights
+            weights = jnp.clip(weights, eps, self.w_bound - eps)
 
-    @transition(output_compartments=["dInputs"])
-    @staticmethod
-    def backtransmit(stride, padding, x_delta_shape, preSpike, postSpike, weights):  ## action-backpropagating routine
+        self.weights.set(weights)
+        self.dWeights.set(dWeights)
+
+    @compilable
+    def backtransmit(self):  ## action-backpropagating routine
         ## calc dInputs
-        dInputs = calc_dX_deconv(weights, postSpike, delta_shape=x_delta_shape,
-                                 stride_size=stride, padding=padding)
-        return dInputs
-
-    @transition(output_compartments=["inputs", "outputs", "preSpike", "postSpike", "preTrace", "postTrace"])
-    @staticmethod
-    def reset(in_shape, out_shape):
-        preVals = jnp.zeros(in_shape)
-        postVals = jnp.zeros(out_shape)
-        inputs = preVals
-        outputs = postVals
-        preSpike = preVals
-        postSpike = postVals
-        preTrace = preVals
-        postTrace = postVals
-        return inputs, outputs, preSpike, postSpike, preTrace, postTrace
+        dInputs = calc_dX_deconv(
+            self.weights.get(), self.postSpike.get(), delta_shape=self.x_delta_shape, stride_size=self.stride,
+            padding=self.padding
+        )
+        self.dInputs.set(dInputs)
+
+    @compilable
+    def reset(self):  # in_shape, out_shape):
+        preVals = jnp.zeros(self.in_shape.get())
+        postVals = jnp.zeros(self.out_shape.get())
+        self.inputs.set(preVals)
+        self.outputs.set(postVals)
+        self.preSpike.set(preVals)
+        self.postSpike.set(postVals)
+        self.preTrace.set(preVals)
+        self.postTrace.set(postVals)
 
     @classmethod
     def help(cls): ## component help function
diff --git a/tests/components/synapses/convolution/test_traceSTDPConvSynapse.py b/tests/components/synapses/convolution/test_traceSTDPConvSynapse.py
@@ -54,14 +54,14 @@ def test_TraceSTDPConvSynapse1():
           [[1.], [0.]]]]
     )
 
-    reset_process.run() #
+    reset_process.run() # ctx.reset()
     a.inputs.set(x)
     advance_process.run(t=1., dt=dt)  # ctx.run(t=1., dt=dt)
     y = (a.outputs.get() > 0.) * 1. ## fake out post-syn spikes
     assert_array_equal(y, y_truth)
-    print(y)
-    print("y.Tr:\n", y_truth)
-    print("======")
+    # print(y)
+    # print("y.Tr:\n", y_truth)
+    # print("======")
 
     # print("NGC-Learn.shape = ", node.outputs.get().shape)
     a.preSpike.set(x)
diff --git a/tests/components/synapses/convolution/test_traceSTDPDeconvSynapse.py b/tests/components/synapses/convolution/test_traceSTDPDeconvSynapse.py
