minor cleanup

Author: Alexander Ororbia

ngclearn/utils/diffeq/ode_utils.py

@@ -127,7 +127,7 @@ def _leapfrog(carry, dfq, dt, params):
     return new_carry, (new_carry, carry)
 
 @partial(jit, static_argnums=(3, 4))
-def leapfrog(t_curr, q_curr, p_curr, dfq, L, step_size, params):
+def leapfrog(t_curr, q_curr, p_curr, dfq, L, step_size, params): ## leapfrog estimator step
     t = t_curr + 0.
     q = q_curr + 0.
     p = p_curr + 0.

ngclearn/utils/model_utils.py

@@ -458,7 +458,8 @@ def d_relu6(x):
     """
     # df/dx = 1 if 0<x<6 else 0
     # I_x = (z >= a_min) *@ (z <= b_max) //create an indicator function  a = 0 b = 6
-    Ix1 = (x > 0.).astype(jnp.float32) #tf.cast(tf.math.greater_equal(x, 0.0),dtype=tf.float32)
+    #Ix1 = (x > 0.).astype(jnp.float32) #tf.cast(tf.math.greater_equal(x, 0.0),dtype=tf.float32)
+    Ix1 = (x >= 0.).astype(jnp.float32)
     Ix2 = (x <= 6.).astype(jnp.float32) #tf.cast(tf.math.less_equal(x, 6.0),dtype=tf.float32)
     Ix = Ix1 * Ix2
     return Ix
@@ -726,34 +727,37 @@ def d_clip(x, min_val, max_val):
     return jnp.where((x < min_val) | (x > max_val), 0.0, 1.0)
 
 
-@partial(jit, static_argnums=[2])
-def lkwta(x, group_masks, nWTA=(20,)): ## local k-WTA
+@partial(jit, static_argnums=[2, 3])
+def lkwta(x, group_masks, nWTA=(20,), clipval=-1.): ## local k-WTA
     out = 0.
     for g in range(len(group_masks)):
         m = group_masks[g]
-        x_g = kwta(x, m, nWTA[g])
+        x_g = kwta(x, m, nWTA[g], clipval)
         out = x_g + out
     return out
 
-@partial(jit, static_argnums=[2])
-def d_lkwta(x, group_masks, nWTA=(20,)): ## d(lkwta(x))/dx
+@partial(jit, static_argnums=[2, 3])
+def d_lkwta(x, group_masks, nWTA=(20,), clipval=-1.): ## d(lkwta(x))/dx
     out = 0.
     for g in range(len(group_masks)):
         m = group_masks[g]
-        x_g = d_kwta(x, m, nWTA[g])
+        x_g = d_kwta(x, m, nWTA[g], clipval)
         out = x_g + out
     return out
 
-@partial(jit, static_argnums=[2])
-def kwta(x, m, nWTA=20):
+@partial(jit, static_argnums=[2, 3])
+def kwta(x, m, nWTA=20, clipval=-1.):
     _x = x * m + (1. - m) * (jnp.amin(x) - 1.)
     values, indices = lax.top_k(_x, nWTA) # Note: we do not care to sort the indices
     kth = jnp.expand_dims(jnp.min(values,axis=1),axis=1) # must do comparison per sample in potential mini-batch
     topK = jnp.greater_equal(_x, kth).astype(jnp.float32) # cast booleans to floats
-    return topK * x
+    topK = topK * x
+    if clipval > 0.:
+        topK = jnp.clip(topK, -clipval, clipval)
+    return topK
 
-@partial(jit, static_argnums=[2])
-def d_kwta(x, m, nWTA=20): ## d(kwta(x))/dx
+@partial(jit, static_argnums=[2, 3])
+def d_kwta(x, m, nWTA=20, clipval=-1.): ## d(kwta(x))/dx
     _x = x * m + (1. - m) * (jnp.amin(x) - 1.)
     values, indices = lax.top_k(_x, nWTA) # Note: we do not care to sort the indices
     kth = jnp.expand_dims(jnp.min(values,axis=1),axis=1) # must do comparison per sample in potential mini-batch

ngclearn/utils/patch_utils.py

@@ -148,7 +148,7 @@ def generate_patch_set(x_batch, patch_size=(8, 8), max_patches=50, center=True,
     for s in range(_x_batch.shape[0]):
         xs = _x_batch[s, :]
         xs = xs.reshape(px, py)
-        patches = extract_patches_2d(xs, patch_size, max_patches=max_patches, random_state=seed, extraction_step=step_size)#, random_state=69)
+        patches = extract_patches_2d(xs, patch_size, max_patches=max_patches, random_state=seed) #, extraction_step=step_size)#, random_state=69)
         patches = np.reshape(patches, (len(patches), -1)) # flatten each patch in set
         if s > 0:
             p_batch = np.concatenate((p_batch,patches),axis=0)

tests/components/synapses/hebbian/test_BCMSynapse.py

@@ -19,7 +19,6 @@ def test_BCMSynapse1():
         a = BCMSynapse(
             name="a", shape=(1,1), tau_w=40., tau_theta=20., key=subkeys[0]
         )
-
         evolve_process = (MethodProcess("evolve_process")
                            >> a.evolve)
 
@@ -42,4 +41,5 @@ def test_BCMSynapse1():
     # print(truth)
     assert_array_equal(a.dWeights.get(), truth)
 
-test_BCMSynapse1()
+#test_BCMSynapse1()
+