Fix missing shape specification

patrick-kidger · Nov 10, 2023 · f666790 · f666790
1 parent 48a27f9
commit f666790
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/docs/tricks.md b/docs/tricks.md
@@ -34,7 +34,7 @@ And here is how to replace the `weight` of every linear layer in some arbitrary
 def trunc_init(weight: jax.Array, key: jax.random.PRNGKey) -> jax.Array:
   out, in_ = weight.shape
   stddev = math.sqrt(1 / in_)
-  return stddev * jax.random.truncated_normal(key, lower=-2, upper=2)
+  return stddev * jax.random.truncated_normal(key, shape=(out, in_), lower=-2, upper=2)
 
 def init_linear_weight(model, init_fn, key):
   is_linear = lambda x: isinstance(x, eqx.nn.Linear)