(inputs, f, c)
| 1442 | |
| 1443 | @subgraph("SyncBnSplitStats", dtype, device, 1) |
| 1444 | def syncbn_split_stats(inputs, f, c): |
| 1445 | stats = inputs[0] |
| 1446 | c_1 = c(1, dtype="int32") |
| 1447 | channel_x1s_end = c(channels + 1, dtype="int32") |
| 1448 | |
| 1449 | def _subtensor(src, axis, begin, end): |
| 1450 | items = ((axis, (begin is not None), (end is not None), False, False),) |
| 1451 | args = () |
| 1452 | if begin is not None: |
| 1453 | args += (begin,) |
| 1454 | if end is not None: |
| 1455 | args += (end,) |
| 1456 | return f(builtin.Subtensor(items=items), src, *args) |
| 1457 | |
| 1458 | reduce_size = _subtensor(stats, 1, None, c_1) |
| 1459 | channel_x1s = _subtensor(stats, 1, c_1, channel_x1s_end) |
| 1460 | channel_x2s = _subtensor(stats, 1, channel_x1s_end, None) |
| 1461 | reduce_size = f(builtin.Reshape(), reduce_size, c_1) |
| 1462 | return (reduce_size, channel_x1s, channel_x2s), (False, True, True) |
| 1463 | |
| 1464 | # fmt : on |
| 1465 | return ( |
no test coverage detected