From 92e51035bfa91340b96a4a8e6bc350d2c0a96091 Mon Sep 17 00:00:00 2001
From: Vineet Bansal <vineetbansal@protonmail.com>
Date: Mon, 27 Jul 2026 15:55:04 -0400
Subject: [PATCH 1/2] Add failing test for edge attributes on a pre-sorted edge
 index

---
 tests/core/test_graph.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tests/core/test_graph.py b/tests/core/test_graph.py
index f27071bb..8dc7dfdd 100644
--- a/tests/core/test_graph.py
+++ b/tests/core/test_graph.py
@@ -35,6 +35,17 @@ def test_init_with_edge_index():
     assert isinstance(g.edge_to_index, dict)
 
 
+@pytest.mark.xfail(reason="edge attributes are indexed with a None permutation", strict=True)
+def test_init_with_presorted_edge_index_keeps_edge_attrs():
+    # An EdgeIndex that already carries sort_order="row" is returned by sort_by without a
+    # permutation, in which case the edge attributes must be left as they are.
+    edge_index = EdgeIndex([[0, 0, 1], [1, 2, 2]], sparse_size=(3, 3), sort_order="row")
+    edge_weight = torch.tensor([1.0, 2.0, 3.0])
+    g = Graph(Data(edge_index=edge_index, num_nodes=3, edge_weight=edge_weight))
+    assert g.data.edge_weight.shape == edge_weight.shape
+    assert torch.equal(g.data.edge_weight, edge_weight)
+
+
 def test_init_with_mapping():
     edge_index = get_random_edge_index(100, 100, 1000)
     data = Data(edge_index=edge_index, num_nodes=100)

From 3406fc9009302731f11523b45a622b97cf25b4f4 Mon Sep 17 00:00:00 2001
From: Vineet Bansal <vineetbansal@protonmail.com>
Date: Mon, 27 Jul 2026 15:56:49 -0400
Subject: [PATCH 2/2] Keep edge attributes intact for an already-sorted edge
 index

---
 src/pathpyG/core/graph.py | 6 ++++--
 tests/core/test_graph.py  | 1 -
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/pathpyG/core/graph.py b/src/pathpyG/core/graph.py
index 52d15470..88950443 100644
--- a/src/pathpyG/core/graph.py
+++ b/src/pathpyG/core/graph.py
@@ -116,8 +116,10 @@ def __init__(self, data: Data, mapping: Optional[IndexMap] = None):
 
         # sort EdgeIndex and validate
         data.edge_index, sorted_idx = data.edge_index.sort_by("row")
-        for edge_attr in self.edge_attrs():
-            data[edge_attr] = self.data[edge_attr][sorted_idx]
+        # an edge index that is already sorted is returned without a permutation
+        if sorted_idx is not None:
+            for edge_attr in self.edge_attrs():
+                data[edge_attr] = self.data[edge_attr][sorted_idx]
 
         data.edge_index.validate()
 
diff --git a/tests/core/test_graph.py b/tests/core/test_graph.py
index 8dc7dfdd..615265ab 100644
--- a/tests/core/test_graph.py
+++ b/tests/core/test_graph.py
@@ -35,7 +35,6 @@ def test_init_with_edge_index():
     assert isinstance(g.edge_to_index, dict)
 
 
-@pytest.mark.xfail(reason="edge attributes are indexed with a None permutation", strict=True)
 def test_init_with_presorted_edge_index_keeps_edge_attrs():
     # An EdgeIndex that already carries sort_order="row" is returned by sort_by without a
     # permutation, in which case the edge attributes must be left as they are.