update model utils

Author: rxng8

ngclearn/utils/model_utils.py

@@ -85,7 +85,7 @@ def create_function(fun_name, args=None):
     if fun_name == "tanh":
         fx = tanh
         dfx = d_tanh
-    elif "bkwta" in fun_name:
+    elif fun_name == "bkwta":
         fx = bkwta
         dfx = bkwta #d_identity
     elif fun_name == "sine":
@@ -651,8 +651,8 @@ def layer_normalize(x, shift=0., scale=1.):
         layer-normalized data samples `x`
     """
     xmu = jnp.mean(x, axis=1, keepdims=True)
-    xsigma = jnp.sqrt(jnp.mean(jnp.square(x - xmu)).clip(min=1e-6))
-    _x = (x - xmu)/(xsigma + 1e-6)
+    xsigma = jnp.sqrt(jnp.mean(jnp.square(x - xmu)).clip(min=1e-6)).clip(min=1e-6)
+    _x = (x - xmu) / xsigma
     return _x * scale + shift
 
 @jit