Fix GCNN graph reductions for IrrepsArray inputs

Author: giotre

src/tensorial/gcnn/graph_ops.py

@@ -1,3 +1,4 @@
+from functools import singledispatch
 from typing import Literal
 
 import e3nn_jax as e3j
@@ -35,6 +36,7 @@ def _prepare_segments(
     return num_segments, segment_ids
 
 
+@singledispatch
 def segment_sum(
     data: Float[jax.Array, "N ..."],
     segment_sizes: Int[jax.Array, "num_segments"],
@@ -85,6 +87,23 @@ def segment_sum(
     return result
 
 
+@segment_sum.register
+def _(
+    data: e3j.IrrepsArray,
+    segment_sizes: Int[jax.Array, "num_segments"],
+    mask: Bool[jax.Array, "N ..."] | None = None,
+    segment_mask: Bool[jax.Array, "num_segments"] | None = None,
+) -> e3j.IrrepsArray:
+    result = segment_sum(
+        data.array,
+        segment_sizes,
+        mask=mask,
+        segment_mask=segment_mask,
+    )
+    return e3j.IrrepsArray(data.irreps, result)
+
+
+@singledispatch
 def segment_mean(
     data: Float[jax.Array, "N ..."],
     segment_sizes: Int[jax.Array, "num_segments"],
@@ -158,6 +177,23 @@ def segment_mean(
     return jnp.where(segment_mask, mean, jnp.zeros_like(mean))
 
 
+@segment_mean.register
+def _(
+    data: e3j.IrrepsArray,
+    segment_sizes: Int[jax.Array, "num_segments"],
+    mask: Bool[jax.Array, "N ..."] | None = None,
+    segment_mask: Bool[jax.Array, "num_segments"] | None = None,
+) -> e3j.IrrepsArray:
+    result = segment_mean(
+        data.array,
+        segment_sizes,
+        mask=mask,
+        segment_mask=segment_mask,
+    )
+    return e3j.IrrepsArray(data.irreps, result)
+
+
+@singledispatch
 def segment_min(
     data: Float[jax.Array, "N ..."],
     segment_sizes: Int[jax.Array, "num_segments"],
@@ -210,6 +246,23 @@ def segment_min(
     return data_min
 
 
+@segment_min.register
+def _(
+    data: e3j.IrrepsArray,
+    segment_sizes: Int[jax.Array, "num_segments"],
+    mask: Bool[jax.Array, "N ..."] | None = None,
+    segment_mask: Bool[jax.Array, "num_segments"] | None = None,
+) -> e3j.IrrepsArray:
+    result = segment_min(
+        data.array,
+        segment_sizes,
+        mask=mask,
+        segment_mask=segment_mask,
+    )
+    return e3j.IrrepsArray(data.irreps, result)
+
+
+@singledispatch
 def segment_max(
     data: Float[jax.Array, "N ..."],
     segment_sizes: Int[jax.Array, "num_segments"],
@@ -262,6 +315,22 @@ def segment_max(
     return data_max
 
 
+@segment_max.register
+def _(
+    data: e3j.IrrepsArray,
+    segment_sizes: Int[jax.Array, "num_segments"],
+    mask: Bool[jax.Array, "N ..."] | None = None,
+    segment_mask: Bool[jax.Array, "num_segments"] | None = None,
+) -> e3j.IrrepsArray:
+    result = segment_max(
+        data.array,
+        segment_sizes,
+        mask=mask,
+        segment_mask=segment_mask,
+    )
+    return e3j.IrrepsArray(data.irreps, result)
+
+
 _REDUCTIONS = {
     "mean": segment_mean,
     "sum": segment_sum,
@@ -276,7 +345,7 @@ def segment_reduce(
     reduction: Literal["sum", "mean", "min", "max"],
     mask: Bool[jax.Array, "N ..."] | None = None,
     segment_mask: Bool[jax.Array, "num_segments"] | None = None,
-) -> Float[jax.Array, "num_segments ..."]:
+) -> Float[jax.Array, "num_segments ..."] | e3j.IrrepsArray:
     """Performs a masked segment reduction over batched graph data.
 
     This function is JAX-jittable and handles the logic for applying a mask
@@ -348,6 +417,6 @@ def _jraph_segment(
     except AttributeError:
         raise ValueError(f"Unknown reduction operation: {reduction}") from None
 
-    return jax.tree_util.tree_map(
+    return jax.tree.map(
         lambda n: op(n, segment_ids, num_segments, indices_are_sorted, unique_indices), inputs
     )

tests/unit/gcnn/test_graph_ops.py

@@ -1,8 +1,9 @@
+import e3nn_jax as e3j
 import jax
 import jax.numpy as jnp
 import pytest
 
-from tensorial.gcnn import graph_ops
+from tensorial.gcnn import graph_ops, keys
 
 
 @pytest.mark.parametrize("jit", [False, True])
@@ -220,3 +221,172 @@ def test_segment_reduce_with_explicit_inf(jit, reduction, expected_inf):
 
     res = op(data, segment_sizes, mask=mask, segment_mask=segment_mask)
     assert jnp.allclose(res, expected)
+
+
+@pytest.mark.parametrize("jit", [False, True])
+def test_segment_sum_irreps_array(jit):
+
+    op = jax.jit(graph_ops.segment_sum) if jit else graph_ops.segment_sum
+
+    x = e3j.IrrepsArray(
+        "2x0e",
+        jnp.array(
+            [
+                [1.0, 2.0],
+                [3.0, 4.0],
+                [5.0, 6.0],
+            ]
+        ),
+    )
+    segment_sizes = jnp.array([2, 1])
+    # Case 1: no segment_mask
+    res = op(x, segment_sizes)
+    assert isinstance(res, e3j.IrrepsArray)
+    assert res.irreps == x.irreps
+    expected = jnp.array(
+        [
+            [4.0, 6.0],
+            [5.0, 6.0],
+        ]
+    )
+    assert jnp.allclose(res.array, expected)
+
+    # Case 2: with segment_mask
+    segment_mask = jnp.array([True, False])
+    res_masked = op(x, segment_sizes, segment_mask=segment_mask)
+    assert isinstance(res_masked, e3j.IrrepsArray)
+    expected_masked = jnp.array(
+        [
+            [4.0, 6.0],
+            [0.0, 0.0],
+        ]
+    )
+    assert jnp.allclose(res_masked.array, expected_masked)
+
+
+@pytest.mark.parametrize("jit", [False, True])
+def test_segment_mean_irreps_array(jit):
+
+    op = jax.jit(graph_ops.segment_mean) if jit else graph_ops.segment_mean
+
+    x = e3j.IrrepsArray(
+        "2x0e",
+        jnp.array(
+            [
+                [1.0, 2.0],
+                [3.0, 4.0],
+                [5.0, 6.0],
+            ]
+        ),
+    )
+    segment_sizes = jnp.array([2, 1])
+
+    # Case 1: no segment_mask
+    res = op(x, segment_sizes)
+    assert isinstance(res, e3j.IrrepsArray)
+    assert res.irreps == x.irreps
+    expected = jnp.array(
+        [
+            [2.0, 3.0],
+            [5.0, 6.0],
+        ]
+    )
+    assert jnp.allclose(res.array, expected)
+
+    # Case 2: with segment_mask and data mask
+    mask = jnp.array([True, False, True])
+    segment_mask = jnp.array([True, False])
+    res_masked = op(x, segment_sizes, mask=mask, segment_mask=segment_mask)
+    assert isinstance(res_masked, e3j.IrrepsArray)
+    expected_masked = jnp.array(
+        [
+            [1.0, 2.0],
+            [0.0, 0.0],
+        ]
+    )
+    assert jnp.allclose(res_masked.array, expected_masked)
+
+
+@pytest.mark.parametrize("jit", [False, True])
+def test_segment_min_max_irreps_array(jit):
+
+    op_min = jax.jit(graph_ops.segment_min) if jit else graph_ops.segment_min
+    op_max = jax.jit(graph_ops.segment_max) if jit else graph_ops.segment_max
+
+    x = e3j.IrrepsArray(
+        "2x0e",
+        jnp.array(
+            [
+                [1.0, 2.0],
+                [3.0, 4.0],
+                [5.0, 6.0],
+            ]
+        ),
+    )
+    segment_sizes = jnp.array([2, 1])
+    mask = jnp.array([True, False, True])
+    segment_mask = jnp.array([True, False])
+
+    # Min test
+    res_min = op_min(x, segment_sizes, mask=mask, segment_mask=segment_mask)
+    assert isinstance(res_min, e3j.IrrepsArray)
+    expected_min = jnp.array(
+        [
+            [1.0, 2.0],
+            [jnp.inf, jnp.inf],
+        ]
+    )
+    assert jnp.allclose(res_min.array, expected_min)
+
+    # Max test
+    res_max = op_max(x, segment_sizes, mask=mask, segment_mask=segment_mask)
+    assert isinstance(res_max, e3j.IrrepsArray)
+    expected_max = jnp.array(
+        [
+            [1.0, 2.0],
+            [-jnp.inf, -jnp.inf],
+        ]
+    )
+    assert jnp.allclose(res_max.array, expected_max)
+
+
+@pytest.mark.parametrize("jit", [False, True])
+def test_graph_segment_reduce_irreps_array_with_node_mask(jit):
+
+    op = (
+        jax.jit(graph_ops.graph_segment_reduce, static_argnums=(1, 2))
+        if jit
+        else graph_ops.graph_segment_reduce
+    )
+
+    x = e3j.IrrepsArray(
+        "2x0e",
+        jnp.array(
+            [
+                [1.0, 2.0],
+                [3.0, 4.0],
+                [5.0, 6.0],
+            ]
+        ),
+    )
+
+    graph = {
+        "nodes": {
+            "features": x,
+            keys.MASK: jnp.array([True, False, True]),
+        },
+        "n_node": jnp.array([2, 1]),
+    }
+
+    res = op(graph, "nodes.features", "mean")
+
+    assert isinstance(res, e3j.IrrepsArray)
+    assert res.irreps == x.irreps
+
+    expected = jnp.array(
+        [
+            [1.0, 2.0],
+            [5.0, 6.0],
+        ]
+    )
+    assert jnp.allclose(res.array, expected)