             switch (arg.out_strides().ndim()) {
                 case 0:
                     reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
-                            MakeReductionKernelArg<In, Out, 3, 0, 3>(arg), reduce_block_size, impl);
-                    return;
-                case 1:
-                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
-                            MakeReductionKernelArg<In, Out, 3, 1, 2>(arg), reduce_block_size, impl);
+                            MakeReductionKernelArg<In, Out, 3, 0, 3>(arg), out_block_size, reduce_block_size, impl);
+                    return;
+                case 1:
+                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
+                            MakeReductionKernelArg<In, Out, 3, 1, 2>(arg), out_block_size, reduce_block_size, impl);
                     return;
                 case 2:
                     reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
-                            MakeReductionKernelArg<In, Out, 3, 2, 1>(arg), reduce_block_size, impl);
+                            MakeReductionKernelArg<In, Out, 3, 2, 1>(arg), out_block_size, reduce_block_size, impl);
                     return;
                 case 3:
                     reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
-                            MakeReductionKernelArg<In, Out, 3, 3, 0>(arg), reduce_block_size, impl);
+                            MakeReductionKernelArg<In, Out, 3, 3, 0>(arg), out_block_size, reduce_block_size, impl);
                     return;
             }
             XCHAINER_NEVER_REACH();

    return x + 1;
}

template <
        typename In,
        typename Out,

            switch (arg.out_strides().ndim()) {
                case 0:
                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
                            MakeReductionKernelArg<In, Out, 1, 0, 1>(arg), out_block_size, reduce_block_size, impl);
                    return;
                case 1:
                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
                            MakeReductionKernelArg<In, Out, 1, 1, 0>(arg), out_block_size, reduce_block_size, impl);
                    return;
            }
            XCHAINER_NEVER_REACH();
