add some utils for ease of use

Author: Jackmin801

tests/test_poly.py

@@ -3,7 +3,7 @@
 import base64
 from toploc.poly import (
     find_injective_modulus,
-    build_proofs,
+    build_proofs_bytes,
     build_proofs_base64,
     ProofPoly,
 )
@@ -132,7 +132,7 @@ def sample_activations():
 
 def test_build_proofs(sample_activations):
     """Test building proofs"""
-    proofs = build_proofs(sample_activations, decode_batching_size=2, topk=5)
+    proofs = build_proofs_bytes(sample_activations, decode_batching_size=2, topk=5)
     assert isinstance(proofs, list)
     assert all(isinstance(p, bytes) for p in proofs)
     assert len(proofs) == 5
@@ -151,7 +151,7 @@ def test_build_proofs_base64(sample_activations):
 
 def test_build_proofs_skip_prefill(sample_activations):
     """Test building proofs with skip_prefill"""
-    proofs = build_proofs(
+    proofs = build_proofs_bytes(
         sample_activations, decode_batching_size=2, topk=5, skip_prefill=True
     )
     assert isinstance(proofs, list)
@@ -165,7 +165,7 @@ def test_build_proofs_error_handling():
         torch.randn(0, 16, dtype=torch.bfloat16),
         torch.randn(16, dtype=torch.bfloat16),
     ]
-    proofs = build_proofs(invalid_activations, decode_batching_size=2, topk=5)
+    proofs = build_proofs_bytes(invalid_activations, decode_batching_size=2, topk=5)
     assert isinstance(proofs, list)
     assert all(isinstance(p, bytes) for p in proofs)
 
@@ -176,21 +176,23 @@ def test_build_proofs_error_handling():
 def test_build_proofs_edge_cases(sample_activations):
     """Test edge cases for proof building"""
     # Test with minimal topk
-    proofs_min = build_proofs(sample_activations, decode_batching_size=2, topk=1)
+    proofs_min = build_proofs_bytes(sample_activations, decode_batching_size=2, topk=1)
     assert len(proofs_min) > 0
 
     # Test with large batching size
-    proofs_large_batch = build_proofs(
+    proofs_large_batch = build_proofs_bytes(
         sample_activations, decode_batching_size=10, topk=5
     )
     assert len(proofs_large_batch) > 0
 
     # Test with only one prefill activation
-    proofs_one = build_proofs(sample_activations[:1], decode_batching_size=2, topk=5)
+    proofs_one = build_proofs_bytes(
+        sample_activations[:1], decode_batching_size=2, topk=5
+    )
     assert len(proofs_one) == 1
 
     # Test with only one activation and skip_prefill
-    proofs_one_skip = build_proofs(
+    proofs_one_skip = build_proofs_bytes(
         sample_activations[:1], decode_batching_size=2, topk=5, skip_prefill=True
     )
     assert len(proofs_one_skip) == 0

toploc/__init__.py

@@ -1,4 +1,4 @@
-from toploc.poly import ProofPoly, build_proofs, build_proofs_base64  # noqa: F401
+from toploc.poly import ProofPoly, build_proofs, build_proofs_bytes, build_proofs_base64  # noqa: F401
 from toploc.utils import sha256sum  # noqa: F401
 
 __version__ = "0.0.0.dev1"

toploc/poly.py

@@ -80,34 +80,6 @@ def build_proofs(
     decode_batching_size: int,
     topk: int,
     skip_prefill: bool = False,
-) -> list[bytes]:
-    return [
-        proof.to_bytes()
-        for proof in _build_proofs(
-            activations, decode_batching_size, topk, skip_prefill
-        )
-    ]
-
-
-def build_proofs_base64(
-    activations: list[torch.Tensor],
-    decode_batching_size: int,
-    topk: int,
-    skip_prefill: bool = False,
-) -> list[str]:
-    return [
-        proof.to_base64()
-        for proof in _build_proofs(
-            activations, decode_batching_size, topk, skip_prefill
-        )
-    ]
-
-
-def _build_proofs(
-    activations: list[torch.Tensor],
-    decode_batching_size: int,
-    topk: int,
-    skip_prefill: bool = False,
 ) -> list[ProofPoly]:
     proofs = []
 
@@ -137,3 +109,47 @@ def _build_proofs(
         )
 
     return proofs
+
+
+def build_proofs_bytes(
+    activations: list[torch.Tensor],
+    decode_batching_size: int,
+    topk: int,
+    skip_prefill: bool = False,
+) -> list[bytes]:
+    return [
+        proof.to_bytes()
+        for proof in build_proofs(activations, decode_batching_size, topk, skip_prefill)
+    ]
+
+
+def build_proofs_base64(
+    activations: list[torch.Tensor],
+    decode_batching_size: int,
+    topk: int,
+    skip_prefill: bool = False,
+) -> list[str]:
+    return [
+        proof.to_base64()
+        for proof in build_proofs(activations, decode_batching_size, topk, skip_prefill)
+    ]
+
+
+def batch_activations(
+    activations: list[torch.Tensor],
+    decode_batching_size: int,
+) -> list[torch.Tensor]:
+    batches = []
+
+    # Prefill
+    flat_view = activations[0].view(-1)
+    batches.append(flat_view)
+
+    # Batched Decode
+    for i in range(1, len(activations), decode_batching_size):
+        flat_view = torch.cat(
+            [i.view(-1) for i in activations[i : i + decode_batching_size]]
+        )
+        batches.append(flat_view)
+
+    return batches