NVIDIA/nvbench
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvbench.git synced 2026-04-19 22:38:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Improvements to readability of examples per PR review

Browse Source
This commit is contained in:
Oleksandr Pavlyk
2025-07-31 16:20:52 -05:00
parent c91204f259
commit 453a1648aa
6 changed files with 30 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/examples/cccl_parallel_segmented_reduce.py
View File

@@ -61,10 +61,6 @@ def segmented_reduce(state: nvbench.State):
dev_id = state.get_device()
cp_stream = as_cp_ExternalStream(state.get_stream(), dev_id)
with cp_stream:
rng = cp.random.default_rng()
mat = rng.integers(low=-31, high=32, dtype=np.int32, size=(n_rows, n_cols))
def add_op(a, b):
return a + b
@@ -84,6 +80,8 @@ def segmented_reduce(state: nvbench.State):
h_init = np.zeros(tuple(), dtype=np.int32)
with cp_stream:
rng = cp.random.default_rng()
mat = rng.integers(low=-31, high=32, dtype=np.int32, size=(n_rows, n_cols))
d_input = mat
d_output = cp.empty(n_rows, dtype=d_input.dtype)