chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,120 @@
+import pickle
+import unittest
+
+import backend as F
+
+import dgl
+import dgl.ndarray as nd
+import numpy as np
+from dgl.frame import Column
+from utils import parametrize_idtype
+
+
+def test_column_subcolumn():
+    data = F.copy_to(
+        F.tensor(
+            [
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 9.0, 0.0],
+                [3.0, 2.0, 1.0, 0.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 4.0, 0.0],
+            ]
+        ),
+        F.ctx(),
+    )
+    original = Column(data)
+
+    # subcolumn from cpu context
+    i1 = F.tensor([0, 2, 1, 3], dtype=F.int64)
+    l1 = original.subcolumn(i1)
+
+    assert len(l1) == i1.shape[0]
+    assert F.array_equal(l1.data, F.gather_row(data, i1))
+
+    # next subcolumn from target context
+    i2 = F.copy_to(F.tensor([0, 2], dtype=F.int64), F.ctx())
+    l2 = l1.subcolumn(i2)
+
+    assert len(l2) == i2.shape[0]
+    i1i2 = F.copy_to(F.gather_row(i1, F.copy_to(i2, F.context(i1))), F.ctx())
+    assert F.array_equal(l2.data, F.gather_row(data, i1i2))
+
+    # next subcolumn also from target context
+    i3 = F.copy_to(F.tensor([1], dtype=F.int64), F.ctx())
+    l3 = l2.subcolumn(i3)
+
+    assert len(l3) == i3.shape[0]
+    i1i2i3 = F.copy_to(
+        F.gather_row(i1i2, F.copy_to(i3, F.context(i1i2))), F.ctx()
+    )
+    assert F.array_equal(l3.data, F.gather_row(data, i1i2i3))
+
+
+def test_serialize_deserialize_plain():
+    data = F.copy_to(
+        F.tensor(
+            [
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 9.0, 0.0],
+                [3.0, 2.0, 1.0, 0.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 4.0, 0.0],
+            ]
+        ),
+        F.ctx(),
+    )
+    original = Column(data)
+
+    serial = pickle.dumps(original)
+    new = pickle.loads(serial)
+    print("new = {}".format(new))
+
+    assert F.array_equal(new.data, original.data)
+
+
+def test_serialize_deserialize_subcolumn():
+    data = F.copy_to(
+        F.tensor(
+            [
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 9.0, 0.0],
+                [3.0, 2.0, 1.0, 0.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 4.0, 0.0],
+            ]
+        ),
+        F.ctx(),
+    )
+    original = Column(data)
+
+    # subcolumn from cpu context
+    i1 = F.tensor([0, 2, 1, 3], dtype=F.int64)
+    l1 = original.subcolumn(i1)
+
+    serial = pickle.dumps(l1)
+    new = pickle.loads(serial)
+
+    assert F.array_equal(new.data, l1.data)
+
+
+def test_serialize_deserialize_dtype():
+    data = F.copy_to(
+        F.tensor(
+            [
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 9.0, 0.0],
+                [3.0, 2.0, 1.0, 0.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [0.0, 2.0, 4.0, 0.0],
+            ]
+        ),
+        F.ctx(),
+    )
+    original = Column(data)
+    original = original.astype(F.int64)
+
+    serial = pickle.dumps(original)
+    new = pickle.loads(serial)
+
+    assert new.dtype == F.int64