(inputs, f, c)
| 1350 | # fmt : off |
| 1351 | @subgraph("SyncBnStage0", dtype, device, 1) |
| 1352 | def syncbn_stage0(inputs, f, c): |
| 1353 | input = inputs[0] |
| 1354 | reduce_shape = c( |
| 1355 | (1, channels) + (1,) * (ndim - 2), dtype="int32", device=device |
| 1356 | ) |
| 1357 | input_shape = f(GetVarShape(), input) |
| 1358 | input_elems = f(Reduce(mode="product", axis=0), input_shape) |
| 1359 | reduce_elems = f(Reduce(mode="product", axis=0), reduce_shape) |
| 1360 | reduce_size = f("//", input_elems, reduce_elems) |
| 1361 | channel_x1s = f(Reduce(mode="sum"), input, reduce_shape) |
| 1362 | channel_x2s = f(Reduce(mode="sum_sqr"), input, reduce_shape) |
| 1363 | reduce_size_f = f(TypeCvt(dtype=dtype), reduce_size) |
| 1364 | return ( |
| 1365 | (reduce_shape, reduce_size_f, channel_x1s, channel_x2s), |
| 1366 | (False, False, True, True), |
| 1367 | ) |
| 1368 | |
| 1369 | @subgraph("SyncBnStage1", dtype, device, 7) |
| 1370 | def syncbn_stage1(inputs, f, c): |
no test coverage detected