| 225 | |
| 226 | |
| 227 | def run_reduce(keep_dims, reduction_name, batch_gen, input_type, output_type=None, layout=None): |
| 228 | batch_fn = batch_gen(input_type) |
| 229 | dali_reduce_fn, numpy_reduce_fn = reduce_fns[reduction_name] |
| 230 | |
| 231 | for axes in batch_fn.valid_axes(): |
| 232 | dali_res_cpu, dali_res_gpu = run_dali( |
| 233 | dali_reduce_fn, |
| 234 | batch_fn, |
| 235 | keep_dims=keep_dims, |
| 236 | axes=axes, |
| 237 | output_type=output_type, |
| 238 | layout=layout, |
| 239 | ) |
| 240 | |
| 241 | batch_fn.reset() |
| 242 | |
| 243 | np_res = run_numpy( |
| 244 | numpy_reduce_fn, batch_fn, keep_dims=keep_dims, axes=axes, output_type=output_type |
| 245 | ) |
| 246 | |
| 247 | for iteration in range(batch_fn.num_iter()): |
| 248 | compare(dali_res_cpu[iteration], np_res[iteration]) |
| 249 | compare(dali_res_gpu[iteration], np_res[iteration]) |
| 250 | |
| 251 | |
| 252 | def test_reduce(): |