| 1368 | |
| 1369 | @subgraph("SyncBnStage1", dtype, device, 7) |
| 1370 | def syncbn_stage1(inputs, f, c): |
| 1371 | input, reduce_size, channel_x1s, channel_x2s, eps = inputs[0:5] |
| 1372 | weight, bias = inputs[5:7] |
| 1373 | channel_mean = f("/", channel_x1s, reduce_size) |
| 1374 | channel_var = f( |
| 1375 | "+", |
| 1376 | f( |
| 1377 | "/", |
| 1378 | f("**", channel_x1s, c(2)), |
| 1379 | f("-", f("*", reduce_size, reduce_size)), |
| 1380 | ), |
| 1381 | f("/", channel_x2s, reduce_size), |
| 1382 | ) |
| 1383 | invsqrt_channel_var = f("**", f(eps_mode, channel_var, eps), c(-0.5)) |
| 1384 | inv_var_wt = f("*", invsqrt_channel_var, weight) |
| 1385 | neg_channel_mean = f("-", channel_mean) |
| 1386 | outvar = f( |
| 1387 | "fma3", |
| 1388 | input, |
| 1389 | inv_var_wt, |
| 1390 | f("+", f("*", neg_channel_mean, inv_var_wt), bias), |
| 1391 | ) |
| 1392 | return (outvar, channel_mean, channel_var), (True, True, True) |
| 1393 | |
| 1394 | @subgraph("SyncBnStage1Inference", dtype, device, 6) |
| 1395 | def syncbn_stage1_inference(inputs, f, c): |