Fixed GroupNorm raising a spurious runtime error (#106)

patrick-kidger · Jun 6, 2022 · b6abc4c · b6abc4c
1 parent aa0d968
commit b6abc4c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/equinox/nn/array_utils.py b/equinox/nn/array_utils.py
@@ -1,7 +1,10 @@
 from typing import Tuple
 
+import jax.numpy as jnp
+
 from ..custom_types import Array
 
 
 def left_broadcast_to(arr: Array, shape: Tuple[int]):
-    return arr.reshape(shape + (1,) * (len(shape) - arr.ndim))
+    arr = arr.reshape(arr.shape + (1,) * (len(shape) - arr.ndim))
+    return jnp.broadcast_to(arr, shape)
diff --git a/tests/test_nn.py b/tests/test_nn.py
@@ -580,6 +580,10 @@ def test_group_norm(getkey):
     x = jrandom.uniform(getkey(), (128,))
     assert gn(x).shape == (128,)
 
+    gn = eqx.nn.GroupNorm(groups=4, channels=128)
+    x = jrandom.uniform(getkey(), (128, 4, 5))
+    assert gn(x).shape == (128, 4, 5)
+
     gn = eqx.nn.GroupNorm(groups=4, channels=128, channelwise_affine=False)
     x = jrandom.uniform(getkey(), (128, 4, 5))
     assert gn(x).shape == (128, 4, 5)