Inital commit

tests/test_planner.py

@@ -0,0 +1,11 @@
+from truecluster.cluster.planner import plan_even_layers
+
+
+def test_even_plan():
+    assignments = plan_even_layers(num_layers=24, max_nodes=3)
+    assert [(a.layer_start, a.layer_end_exclusive) for a in assignments] == [(0, 8), (8, 16), (16, 24)]
+
+
+def test_remainder_plan():
+    assignments = plan_even_layers(num_layers=25, max_nodes=4)
+    assert [(a.layer_start, a.layer_end_exclusive) for a in assignments] == [(0, 7), (7, 13), (13, 19), (19, 25)]

tests/test_tensors.py

@@ -0,0 +1,11 @@
+import torch
+
+from truecluster.protocol.tensors import deserialize_tensor, serialize_tensor
+
+
+def test_tensor_roundtrip_float16():
+    x = torch.randn(2, 3).half()
+    y = deserialize_tensor(serialize_tensor(x))
+    assert y.dtype == torch.float16
+    assert y.shape == x.shape
+    assert torch.equal(x, y)