Update element_wise.arrangement to flatten the tensors before tiling

voltjia · voltjia · commit a25d837587e9 · 2025-07-11T15:03:50.000+08:00
diff --git a/src/ntops/kernels/element_wise.py b/src/ntops/kernels/element_wise.py
@@ -9,8 +9,7 @@ def arrangement(*tensors, block_size=None):
 
     assert all(tensor.ndim == ndim or tensor.ndim == 0 for tensor in tensors)
 
-    block_shape = tuple(1 for _ in range(ndim - 1)) + (block_size,)
-
     return tuple(
-        tensor.tile(block_shape) if tensor.ndim != 0 else tensor for tensor in tensors
+        tensor.flatten().tile((block_size,)) if tensor.ndim != 0 else tensor
+        for tensor in tensors
     )

Original file line number	Diff line number	Diff line change
`@@ -9,8 +9,7 @@ def arrangement(*tensors, block_size=None):`
`9`	`9`
`10`	`10`	`assert all(tensor.ndim == ndim or tensor.ndim == 0 for tensor in tensors)`
`11`	`11`
`12`		`- block_shape = tuple(1 for _ in range(ndim - 1)) + (block_size,)`
`13`		`-`
`14`	`12`	`return tuple(`
`15`		`- tensor.tile(block_shape) if tensor.ndim != 0 else tensor for tensor in tensors`
	`13`	`+ tensor.flatten().tile((block_size,)) if tensor.ndim != 0 else tensor`
	`14`	`+ for tensor in tensors`
`16`	`15`	`)`