cleaned up gauss/laplace error cells

Author: Alexander Ororbia

ngclearn/components/neurons/graded/gaussianErrorCell.py

@@ -81,7 +81,7 @@ def _advance_state(dt, mu, target, Sigma, modulator, mask): ## compute Gaussian
         dmu = dmu * modulator * mask ## not sure how mask will apply to a full covariance...
         dtarget = dtarget * modulator * mask
         mask = mask * 0. + 1. ## "eat" the mask as it should only apply at time t
-        return dmu, dtarget, dSigma, L[0, 0], mask
+        return dmu, dtarget, dSigma, jnp.squeeze(L), mask
 
     @resolver(_advance_state)
     def advance_state(self, dmu, dtarget, dSigma, L, mask):

ngclearn/components/neurons/graded/laplacianErrorCell.py

@@ -43,7 +43,7 @@ def __init__(self, name, n_units, batch_size=1, scale=1., shape=None, **kwargs):
         else:
             _shape = (batch_size, shape[0], shape[1], shape[2])  ## shape is 4D tensor
         scale_shape = (1, 1)
-        if not isinstance(scale, float):
+        if not isinstance(scale, float) and not isinstance(sigma, int):
             scale_shape = jnp.array(scale).shape
         self.scale_shape = scale_shape
         ## Layer Size setup
@@ -83,7 +83,7 @@ def _advance_state(dt, shift, target, Scale, modulator, mask): ## compute Laplac
         dshift = dshift * modulator * mask
         dtarget = dtarget * modulator * mask
         mask = mask * 0. + 1.  ## "eat" the mask as it should only apply at time t
-        return dshift, dtarget, dScale, L[0, 0], mask
+        return dshift, dtarget, dScale, jnp.squeeze(L), mask
 
     @resolver(_advance_state)
     def advance_state(self, dshift, dtarget, dScale, L, mask):