add group size assertions

Author: liuanji

src/pyjuice/layer/prod_layer.py

@@ -412,7 +412,7 @@ def _forward_backward(self, node_vals: torch.Tensor, element_vals: torch.Tensor,
         accum = 1 if accum else 0
         partial_eval = 1 if local_ids is not None else 0
 
-        assert num_edges & (num_edges - 1) == 0, "`num_edges` must be power of 2."
+        assert num_edges & (num_edges - 1) == 0, "`num_edges` must be a power of 2."
 
         # Fall back to the `torch.compile` kernel in the case where we cannot store child edges within a single block
         if num_edges > 1024:

src/pyjuice/model/backend/par_update.py

@@ -47,7 +47,7 @@ def _record_par_blks(par_start_ids, pflow_start_ids, blk_sizes, blk_intervals, g
 @torch.no_grad()
 def compile_par_update_fn(root_ns: CircuitNodes, BLOCK_SIZE: int = 32, buffer_inc_interval: int = 10000, use_numba: bool = True):
 
-    assert BLOCK_SIZE & (BLOCK_SIZE - 1) == 0, "`BLOCK_SIZE` must be power of 2."
+    assert BLOCK_SIZE & (BLOCK_SIZE - 1) == 0, "`BLOCK_SIZE` must be a power of 2."
 
     par_start_ids = np.zeros([buffer_inc_interval], dtype = np.int64)
     pflow_start_ids = np.zeros([buffer_inc_interval], dtype = np.int64)

src/pyjuice/nodes/construction.py

@@ -21,6 +21,8 @@
 def inputs(var: Union[int,Sequence[int]], num_node_groups: int = 0, dist: Distribution = Distribution(), 
            params: Optional[Tensor] = None, num_nodes: int = 0, group_size: int = 0, **kwargs):
 
+    assert group_size == 0 or group_size & (group_size - 1) == 0, "`group_size` must be a power of 2."
+
     if num_nodes > 0:
         assert num_node_groups == 0, "Only one of `num_nodes` and `num_node_groups` can be set at the same time."
         if group_size == 0:
@@ -72,6 +74,8 @@ def multiply(nodes1: ProdNodesChs, *args, edge_ids: Optional[Tensor] = None, spa
 def summate(nodes1: SumNodesChs, *args, num_nodes: int = 0, num_node_groups: int = 0, 
             edge_ids: Optional[Tensor] = None, group_size: int = 0, **kwargs):
 
+    assert group_size == 0 or group_size & (group_size - 1) == 0, "`group_size` must be a power of 2."
+
     if num_nodes > 0:
         assert num_node_groups == 0, "Only one of `num_nodes` and `num_node_groups` can be set at the same time."
         if group_size == 0:
@@ -101,6 +105,8 @@ def summate(nodes1: SumNodesChs, *args, num_nodes: int = 0, num_node_groups: int
 class set_group_size(_DecoratorContextManager):
     def __init__(self, group_size: int = 1):
 
+        assert group_size & (group_size - 1) == 0, "`group_size` must be a power of 2."
+
         self.group_size = group_size
 
         self.original_group_size = None