diff --git a/tinygrad/utils.py b/tinygrad/utils.py
index 31791b90eb..2887748d62 100644
--- a/tinygrad/utils.py
+++ b/tinygrad/utils.py
@@ -1,7 +1,7 @@
 import numpy as np
 
-def layer_init_uniform(m, h):
-  ret = np.random.uniform(-1., 1., size=(m,h))/np.sqrt(m*h)
+def layer_init_uniform(*x):
+  ret = np.random.uniform(-1., 1., size=x)/np.sqrt(np.prod(x))
   return ret.astype(np.float32)
 
 def fetch_mnist():