revised tests to no longer use weight_distribution/revisions throughout

Author: Alexander Ororbia

ngclearn/components/synapses/patched/hebbianPatchedSynapse.py

@@ -70,7 +70,7 @@ def _calc_update(
 
     dW = dW + prior_lmbda * dW_reg
 
-    if mask!=None:
+    if mask != None:
         dW = dW * mask
 
     return dW * signVal, db * signVal
@@ -95,12 +95,12 @@ def _enforce_constraints(W, block_mask, w_bound, is_nonnegative=True):
     """
     _W = W
     if w_bound > 0.:
-        if is_nonnegative == True:
+        if is_nonnegative:
             _W = jnp.clip(_W, 0., w_bound)
         else:
             _W = jnp.clip(_W, -w_bound, w_bound)
 
-    if block_mask!=None:
+    if block_mask != None:
         _W = _W * block_mask
 
     return _W

ngclearn/components/synapses/patched/patchedSynapse.py

@@ -27,16 +27,16 @@ def _create_multi_patch_synapses(key, shape, n_sub_models, sub_stride, weight_in
 
         shape_ = (end_i - start_i, end_j - start_j) # (di + 2 * si, dj + 2 * sj)
 
-        weights[start_i: end_i,
-                start_j: end_j] = weight_init(shape_, key[2])
+        ## 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[:si,:] = 0.
-        weights[-si:,:] = 0.
-    if sj!=0:
-        weights[:,:sj] = 0.
-        weights[:, -sj:] = 0.
+    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
 
@@ -109,7 +109,8 @@ 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 = {"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,

tests/components/synapses/convolution/test_hebbianConvSynapse.py

@@ -3,7 +3,7 @@
 np.random.seed(42)
 
 from ngclearn import Context, MethodProcess
-import ngclearn.utils.weight_distribution as dist
+from ngclearn.utils.distribution_generator import DistributionGenerator as dist
 from ngclearn.components.synapses.convolution.hebbianConvSynapse import HebbianConvSynapse
 from numpy.testing import assert_array_equal
 

tests/components/synapses/convolution/test_hebbianDeconvSynapse.py

@@ -3,7 +3,7 @@
 np.random.seed(42)
 
 from ngclearn import Context, MethodProcess
-import ngclearn.utils.weight_distribution as dist
+from ngclearn.utils.distribution_generator import DistributionGenerator as dist
 from ngclearn.components.synapses.convolution.hebbianDeconvSynapse import HebbianDeconvSynapse
 from numpy.testing import assert_array_equal
 

tests/components/synapses/convolution/test_traceSTDPConvSynapse.py

@@ -4,7 +4,7 @@
 np.random.seed(42)
 
 from ngclearn import Context, MethodProcess
-import ngclearn.utils.weight_distribution as dist
+from ngclearn.utils.distribution_generator import DistributionGenerator as dist
 from ngclearn.components.synapses.convolution.traceSTDPConvSynapse import TraceSTDPConvSynapse
 from numpy.testing import assert_array_equal
 

tests/components/synapses/convolution/test_traceSTDPDeconvSynapse.py

@@ -4,7 +4,7 @@
 np.random.seed(42)
 
 from ngclearn import Context, MethodProcess
-import ngclearn.utils.weight_distribution as dist
+from ngclearn.utils.distribution_generator import DistributionGenerator as dist
 from ngclearn.components.synapses.convolution.traceSTDPDeconvSynapse import TraceSTDPDeconvSynapse
 from numpy.testing import assert_array_equal
 

tests/components/synapses/hebbian/test_BCMSynapse.py

@@ -4,7 +4,7 @@
 np.random.seed(42)
 
 from ngclearn import Context, MethodProcess
-import ngclearn.utils.weight_distribution as dist
+#from ngclearn.utils.distribution_generator import DistributionGenerator as dist
 from ngclearn.components.synapses.hebbian.BCMSynapse import BCMSynapse
 from numpy.testing import assert_array_equal
 

tests/components/synapses/hebbian/test_eventSTDPSynapse.py

@@ -4,7 +4,7 @@
 np.random.seed(42)
 
 from ngclearn import Context, MethodProcess
-import ngclearn.utils.weight_distribution as dist
+#from ngclearn.utils.distribution_generator import DistributionGenerator as dist
 from ngclearn.components.synapses.hebbian.eventSTDPSynapse import EventSTDPSynapse
 from numpy.testing import assert_array_equal
 

tests/components/synapses/hebbian/test_expSTDPSynapse.py

@@ -5,7 +5,7 @@
 np.random.seed(42)
 
 from ngclearn import Context, MethodProcess
-import ngclearn.utils.weight_distribution as dist
+#from ngclearn.utils.distribution_generator import DistributionGenerator as dist
 from ngclearn.components.synapses.hebbian.expSTDPSynapse import ExpSTDPSynapse
 from numpy.testing import assert_array_equal
 

tests/components/synapses/hebbian/test_traceSTDPSynapse.py

@@ -4,7 +4,7 @@
 np.random.seed(42)
 
 from ngclearn import Context, MethodProcess
-#import ngclearn.utils.weight_distribution as dist
+#from ngclearn.utils.distribution_generator import DistributionGenerator as dist
 from ngclearn.components.synapses.hebbian.traceSTDPSynapse import TraceSTDPSynapse
 from numpy.testing import assert_array_equal