fix: honor zero thresholds when realizing agent graphs

swarm/optimizer/edge_optimizer/parameterization.py

@@ -105,7 +105,7 @@ def realize(self,
             addable_if_not_used_learned_order = (not use_learned_order) and (not _graph.check_cycle(in_node, {out_node}, set()))
             if addable_if_not_used_learned_order or addable_if_use_learned_order:
                 edge_prob = torch.sigmoid(edge_logit / temperature)
-                if threshold:
+                if threshold is not None:
                     edge_prob = torch.tensor(1 if edge_prob > threshold else 0)
                 if torch.rand(1) < edge_prob:
                     out_node.add_successor(in_node)

test/swarm/optimizer/test_edge_threshold.py

@@ -0,0 +1,37 @@
+import pytest
+import torch
+
+from swarm.graph.composite_graph import CompositeGraph
+from swarm.graph.node import Node
+from swarm.optimizer.edge_optimizer.parameterization import EdgeWiseDistribution
+
+pytestmark = pytest.mark.mock_llm
+
+
+class EmptyNode(Node):
+    async def _execute(self, inputs, **kwargs):
+        return {"output": "test"}
+
+
+@pytest.mark.parametrize("learned_order", [False, True])
+@pytest.mark.parametrize("threshold, expected_edges", [(0.0, 1), (1.0, 0), (None, 0)])
+def test_threshold_zero_is_deterministic(monkeypatch, learned_order, threshold, expected_edges):
+    source = EmptyNode("source", "source", True)
+    target = EmptyNode("target", "target", True)
+    graph = CompositeGraph(target, "test")
+    graph.add_node(source)
+    distribution = EdgeWiseDistribution([("source", "target")], initial_probability=0.25)
+    monkeypatch.setattr(torch, "rand", lambda *args: torch.tensor([0.9]))
+    monkeypatch.setattr(
+        distribution,
+        "realize_ranks",
+        lambda *args: ({"source": 0, "target": 1}, torch.tensor(0.0)),
+    )
+
+    realized, log_prob = distribution.realize(
+        graph, threshold=threshold, use_learned_order=learned_order
+    )
+
+    assert realized.num_edges == expected_edges
+    assert torch.isfinite(log_prob)
+    assert graph.num_edges == 0