(inputs, f, c)
| 28 | gopt_level=gopt_level, |
| 29 | ) |
| 30 | def batch_norm_nd(inputs, f, c): |
| 31 | input, eps, weight, bias = inputs[0:4] |
| 32 | reduce_shape = c( |
| 33 | (1, channels) + (1,) * (ndim - 2), dtype="int32", device=device |
| 34 | ) |
| 35 | input_shape = f(GetVarShape(), input) |
| 36 | input_elems = f(Reduce(mode="product", axis=0), input_shape) |
| 37 | reduce_elems = f(Reduce(mode="product", axis=0), reduce_shape) |
| 38 | reduce_size = f("//", input_elems, reduce_elems) |
| 39 | reduce_size = f(TypeCvt(dtype=dtype), reduce_size) |
| 40 | channel_x1s = f(Reduce(mode="sum"), input, reduce_shape) |
| 41 | channel_x2s = f(Reduce(mode="sum_sqr"), input, reduce_shape) |
| 42 | channel_mean = f("/", channel_x1s, reduce_size) |
| 43 | channel_var = f( |
| 44 | "-", f("/", channel_x2s, reduce_size), f("*", channel_mean, channel_mean), |
| 45 | ) |
| 46 | invsqrt_channel_var = f("**", f("+", channel_var, eps), c(-0.5)) |
| 47 | inv_var_wt = f("*", invsqrt_channel_var, weight) |
| 48 | neg_channel_mean = f("-", channel_mean) |
| 49 | outvar = f( |
| 50 | "fma3", input, inv_var_wt, f("fma3", neg_channel_mean, inv_var_wt, bias), |
| 51 | ) |
| 52 | return (outvar,), (True,) |
| 53 | |
| 54 | return batch_norm_nd |
| 55 |
nothing calls this directly
no test coverage detected