Refactor patch synapse (#138)

* Refactor multi-patch synapse creation and initialization

Refactor _create_multi_patch_synapses function to use n_modules instead of n_sub_models and update weight initialization. Introduce weight masks for synaptic weights.

* Refactor HebbianPatchedSynapse and add attributes

Refactor HebbianPatchedSynapse initialization and add new attributes for post-in and pre-out.

Author: Faezehabibi

ngclearn/components/synapses/patched/hebbianPatchedSynapse.py

@@ -197,8 +197,8 @@ def __init__(
             pre_wght=1., post_wght=1., p_conn=1., resist_scale=1., batch_size=1, **kwargs
     ):
         super().__init__(
-            name, shape, n_sub_models, stride_shape, block_mask, weight_init, bias_init, resist_scale, p_conn,
-            batch_size=batch_size, **kwargs
+            name, shape, n_sub_models, stride_shape, weight_init, bias_init, resist_scale, p_conn,
+            batch_size, **kwargs
         )
 
         prior_type, prior_lmbda = prior
@@ -288,6 +288,8 @@ def reset(self):
         # NOTE: Quick workaround is to check if targeted is in the input or not
         hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(preVals) # inputs
         self.outputs.set(postVals) # outputs
+        self.post_in.set(postVals) # post_in
+        self.pre_out.set(preVals) # pre_out
         self.pre.set(preVals) # pre
         self.post.set(postVals) # post
         self.dWeights.set(jnp.zeros(self.shape)) # dW
@@ -304,6 +306,7 @@ def help(cls): ## component help function
         compartment_props = {
             "inputs":
                 {"inputs": "Takes in external input signal values",
+                 "post_in": "Takes in external input signal values",
                  "pre": "Pre-synaptic statistic for Hebb rule (z_j)",
                  "post": "Post-synaptic statistic for Hebb rule (z_i)"},
             "states":
@@ -314,7 +317,8 @@ def help(cls): ## component help function
                 {"dWeights": "Synaptic weight value adjustment matrix produced at time t",
                  "dBiases": "Synaptic bias/base-rate value adjustment vector produced at time t"},
             "outputs":
-                {"outputs": "Output of synaptic transformation"},
+                {"outputs": "Output of synaptic transformation",
+                 "pre_out": "Output of synaptic transformation"},
         }
         hyperparams = {
             "shape": "Overall shape of synaptic weight value matrix; number inputs x number outputs",
@@ -351,3 +355,10 @@ def help(cls): ## component help function
     plt.imshow(Wab.weights.get(), cmap='gray')
     plt.show()
 
+
+
+
+
+
+
+

ngclearn/components/synapses/patched/patchedSynapse.py

@@ -4,78 +4,55 @@
 from jax import random, numpy as jnp, jit
 from ngclearn.components.jaxComponent import JaxComponent
 from ngclearn.utils.distribution_generator import DistributionGenerator
-
 from ngcsimlib.logger import info
 from ngclearn import compilable #from ngcsimlib.parser import compilable
 from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
-# from ngclearn.utils.weight_distribution import initialize_params
-
-
-# def _create_multi_patch_synapses(key, shape, n_sub_models, sub_stride, weight_init):
-#     sub_shape = (shape[0] // n_sub_models, shape[1] // n_sub_models)
-#     di, dj = sub_shape
-#     si, sj = sub_stride
-
-#     weight_shape = ((n_sub_models * di) + 2 * si, (n_sub_models * dj) + 2 * sj)
-#     #weights = initialize_params(key[2], {"dist": "constant", "value": 0.}, weight_shape, use_numpy=True)
-#     large_weight_init = DistributionGenerator.constant(value=0.)
-#     weights = large_weight_init(weight_shape, key[2])
-
-#     for i in range(n_sub_models):
-#         start_i = i * di
-#         end_i = (i + 1) * di + 2 * si
-#         start_j = i * dj
-#         end_j = (i + 1) * dj + 2 * sj
-
-#         shape_ = (end_i - start_i, end_j - start_j) # (di + 2 * si, dj + 2 * sj)
-
-#         ## FIXME: this line below might be wonky...
-#         weights.at[start_i: end_i, start_j: end_j].set( weight_init(shape_, key[2]) )
-#         # weights[start_i : end_i,
-#         #         start_j : end_j] = initialize_params(key[2], init_kernel=weight_init, shape=shape_, use_numpy=True)
-#     if si != 0:
-#         weights.at[:si,:].set(0.) ## FIXME: this setter line might be wonky...
-#         weights.at[-si:,:].set(0.) ## FIXME: this setter line might be wonky...
-#     if sj != 0:
-#         weights.at[:,:sj].set(0.) ## FIXME: this setter line might be wonky...
-#         weights.at[:, -sj:].set(0.) ## FIXME: this setter line might be wonky...
-
-#     return weights
-
-def _create_multi_patch_synapses(key, shape, n_sub_models, sub_stride, weight_init):
-    sub_shape = (shape[0] // n_sub_models, shape[1] // n_sub_models)
-    di, dj = sub_shape
-    si, sj = sub_stride
-
-    weight_shape = ((n_sub_models * di) + 2 * si, (n_sub_models * dj) + 2 * sj)
-    # weights = initialize_params(key[2], {"dist": "constant", "value": 0.}, weight_shape, use_numpy=True)
-    weights = DistributionGenerator.constant(value=0.)(weight_shape, key[2])
-
-    for i in range(n_sub_models):
+
+def _create_multi_patch_synapses(key, shape, n_modules, module_stride=(0, 0), initialization_type=DistributionGenerator.fan_in_gaussian()):
+    key, *subkey = random.split(key, n_modules+10)
+
+    module_shape = (shape[0] // n_modules, shape[1] // n_modules)
+    di, dj = module_shape
+    si, sj = module_stride
+
+    module_shape = di + (2 * si), dj + (2 * sj)
+
+
+    weight_shape = ((n_modules * di) + 2 * si, (n_modules * dj) + 2 * sj)
+    weights = jnp.zeros(weight_shape)
+    w_masks = jnp.zeros(weight_shape)
+
+    for i in range(n_modules):
         start_i = i * di
         end_i = (i + 1) * di + 2 * si
         start_j = i * dj
         end_j = (i + 1) * dj + 2 * sj
 
-        shape_ = (end_i - start_i, end_j - start_j) # (di + 2 * si, dj + 2 * sj)
+        shape_ = (end_i - start_i, end_j - start_j)  # (di + 2 * si, dj + 2 * sj)
 
-        # weights[start_i : end_i,
-        #         start_j : end_j] = initialize_params(key[2],
-        #                                              init_kernel=weight_init,
-        #                                              shape=shape_,
-        #                                              use_numpy=True)
         weights = weights.at[start_i : end_i,
-                start_j : end_j].set(weight_init(shape_, key[2]))
+                start_j : end_j].set(initialization_type(shape_, subkey[i]))
+
+        w_masks = w_masks.at[start_i : end_i,
+                start_j : end_j].set(jnp.ones(shape_))
+
     if si!=0:
         weights = weights.at[:si,:].set(0.)
         weights = weights.at[-si:,:].set(0.)
+
+        w_masks = w_masks.at[:si,:].set(0.)
+        w_masks = w_masks.at[-si:,:].set(0.)
+
     if sj!=0:
         weights = weights.at[:,:sj].set(0.)
         weights = weights.at[:, -sj:].set(0.)
 
-    return weights
+        w_masks = weights.at[:,:sj].set(0.)
+        w_masks = weights.at[:, -sj:].set(0.)
 
 
+    return weights, module_shape, w_masks
+
 class PatchedSynapse(JaxComponent): ## base patched synaptic cable
     """
     A patched dense synaptic cables that creates multiple small dense synaptic cables; no form of synaptic evolution/adaptation
@@ -114,7 +91,7 @@ class PatchedSynapse(JaxComponent): ## base patched synaptic cable
         bias_init: a kernel to drive initialization of biases for this synaptic cable
             (Default: None, which turns off/disables biases)
 
-        block_mask: weight mask matrix
+        w_masks: weight mask matrix
 
         pre_wght: pre-synaptic weighting factor (Default: 1.)
 
@@ -127,8 +104,7 @@ class PatchedSynapse(JaxComponent): ## base patched synaptic cable
             this to < 1. will result in a sparser synaptic structure
     """
 
-    def __init__(
-            self, name, shape, n_sub_models=1, stride_shape=(0,0), block_mask=None, weight_init=None, bias_init=None,
+    def __init__(self, name, shape, n_sub_models=1, stride_shape=(0,0), weight_init=None, bias_init=None,
             resist_scale=1., p_conn=1., batch_size=1, **kwargs
     ):
         super().__init__(name, **kwargs)
@@ -144,60 +120,63 @@ def __init__(
         tmp_key, *subkeys = random.split(self.key.get(), 4)
         if self.weight_init is None:
             info(self.name, "is using default weight initializer!")
-            #self.weight_init = {"dist": "fan_in_gaussian"}
             self.weight_init = DistributionGenerator.fan_in_gaussian()
 
-        weights = _create_multi_patch_synapses(
-            key=subkeys, shape=shape, n_sub_models=self.n_sub_models, sub_stride=self.sub_stride,
-            weight_init=self.weight_init
-        )
-
-        self.block_mask = jnp.where(weights!=0, 1, 0)
-        self.sub_shape = (shape[0]//n_sub_models, shape[1]//n_sub_models)
+        weights, self.sub_shape, self.w_masks = _create_multi_patch_synapses(
+            key=tmp_key, shape=shape, n_modules=self.n_sub_models, module_stride=self.sub_stride,
+            initialization_type = self.weight_init
+                                                                             )
 
         self.shape = weights.shape
-        self.sub_shape = self.sub_shape[0]+(2*self.sub_stride[0]), self.sub_shape[1]+(2*self.sub_stride[1])
-
+        
         if 0. < p_conn < 1.: ## only non-zero and <1 probs allowed
             mask = random.bernoulli(subkeys[1], p=p_conn, shape=self.shape)
             weights = weights * mask ## sparsify matrix
 
         ## Compartment setup
         preVals = jnp.zeros((self.batch_size, self.shape[0]))
         postVals = jnp.zeros((self.batch_size, self.shape[1]))
+
         self.inputs = Compartment(preVals)
         self.outputs = Compartment(postVals)
         self.weights = Compartment(weights)
 
+        self.post_in = Compartment(postVals)
+        self.pre_out = Compartment(preVals)
+        self.weights_T = Compartment(weights.T)
+
         ## Set up (optional) bias values
         if self.bias_init is None:
             info(self.name, "is using default bias value of zero (no bias "
                             "kernel provided)!")
         self.biases = Compartment(self.bias_init((1, self.shape[1]), subkeys[2]) if bias_init else 0.0)
-        #elf.biases = Compartment(initialize_params(subkeys[2], bias_init, (1, self.shape[1])) if bias_init else 0.0)
 
     @compilable
     def advance_state(self):
         # Get the variables
         inputs = self.inputs.get()
+        post_in = self.post_in.get()
         weights = self.weights.get()
         biases = self.biases.get()
 
         outputs = (jnp.matmul(inputs, weights) * self.Rscale) + biases
+        pre_out = jnp.matmul(post_in, weights.T)
 
         # Update compartment
         self.outputs.set(outputs)
+        self.pre_out.set(pre_out)
 
     @compilable
     def reset(self):
         preVals = jnp.zeros((self.batch_size, self.shape[0]))
         postVals = jnp.zeros((self.batch_size, self.shape[1]))
-        inputs = preVals
-        outputs = postVals
+        
         # BUG: the self.inputs here does not have the targeted field
         # NOTE: Quick workaround is to check if targeted is in the input or not
-        hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(inputs)
-        self.outputs.set(outputs)
+        hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(preVals)
+        self.outputs.set(postVals)
+        self.post_in.set(postVals)
+        self.pre_out.set(preVals)
 
     @classmethod
     def help(cls): ## component help function
@@ -208,13 +187,15 @@ def help(cls): ## component help function
         }
         compartment_props = {
             "inputs":
-                {"inputs": "Takes in external input signal values"},
+                {"inputs": "Takes in external input signal values",
+                 "post_in": "Takes in external input signal values"},
             "states":
                 {"weights": "Synapse efficacy/strength parameter values",
                  "biases": "Base-rate/bias parameter values",
                  "key": "JAX PRNG key"},
             "outputs":
-                {"outputs": "Output of synaptic transformation"},
+                {"outputs": "Output of synaptic transformation",
+                 "pre_out": "Output of synaptic transformation"},
         }
         hyperparams = {
             "shape": "Overall shape of synaptic weight value matrix; number inputs x number outputs",
@@ -224,7 +205,7 @@ def help(cls): ## component help function
             "weight_init": "Initialization conditions for synaptic weight (W) values",
             "bias_init": "Initialization conditions for bias/base-rate (b) values",
             "resist_scale": "Resistance level scaling factor (Rscale); applied to output of transformation",
-            "block_mask": "weight mask matrix",
+            "w_masks": "weight mask matrix",
             "p_conn": "Probability of a connection existing (otherwise, it is masked to zero)"
         }
         info = {cls.__name__: properties,
@@ -241,3 +222,8 @@ def help(cls): ## component help function
     plt.imshow(Wab.weights.get(), cmap='gray')
     plt.show()
 
+
+
+
+
+