add special case for scatter_mean

e3nn_jax/_src/scatter.py

@@ -97,6 +97,30 @@ def scatter_mean(
     Returns:
         `jax.numpy.ndarray` or `IrrepsArray`: output array of shape ``(output_size, ...)``
     """
+    if map_back and nel is not None:
+        assert dst is None
+        assert output_size is None
+
+        total = _scatter_op(
+            "sum",
+            0.0,
+            data,
+            nel=nel,
+            map_back=False,
+            mode=mode,
+        )
+        den = jnp.maximum(1, nel)
+
+        for _ in range(total.ndim - nel.ndim):
+            den = den[..., None]
+
+        output = total / den.astype(total.dtype)
+        output = jax.tree_map(
+            lambda x: jnp.repeat(x, nel, axis=0, total_repeat_length=data.shape[0]),
+            output,
+        )
+        return output
+
     total = _scatter_op(
         "sum",
         0.0,

e3nn_jax/_src/scatter_test.py

@@ -71,6 +71,34 @@ def test_scatter_mean():
     )
 
 
+def test_scatter_mean_irreps_array():
+    x = e3nn.IrrepsArray(
+        "0e", jnp.array([[[2.0], [3.0]], [[0.0], [3.0]], [[-10.0], [42.0]]])
+    )
+    dst = jnp.array([[0, 2], [2, 2], [0, 1]])
+
+    np.testing.assert_allclose(  # dst
+        e3nn.scatter_mean(x, dst=dst, output_size=3).array,
+        jnp.array([-4.0, 42.0, 2.0])[..., None],
+    )
+
+    np.testing.assert_allclose(  # map_back
+        e3nn.scatter_mean(x, dst=dst, map_back=True).array,
+        jnp.array([[-4.0, 2.0], [2.0, 2.0], [-4.0, 42.0]])[..., None],
+    )
+
+    x = e3nn.IrrepsArray("0e", jnp.array([[10.0], [1.0], [2.0], [3.0]]))
+    nel = jnp.array([1, 0, 3])
+    np.testing.assert_allclose(  # nel
+        e3nn.scatter_mean(x, nel=nel).array, jnp.array([10.0, 0.0, 2.0])[..., None]
+    )
+
+    np.testing.assert_allclose(  # nel + map_back
+        e3nn.scatter_mean(x, nel=nel, map_back=True).array,
+        jnp.array([10.0, 2.0, 2.0, 2.0])[..., None],
+    )
+
+
 def test_scatter_max():
     jax.config.update("jax_debug_infs", False)