let shape = [1,-1,1,1]; # line up shapes. mean = nn.reshape(mean, shape) variance = nn.reshape(variance, shape) scale = nn.reshape(scale, shape)

let shape = [1,-1,1,1]; # line up shapes. mean = nn.reshape(mean, shape) variance = nn.reshape(variance, shape) scale = nn.reshape(scale, shape) bias = nn.reshape(bias, shape) # calculate 1/sqrt(variance); add an epsilon to prevent dividing by 0 epsilon = nn.constant(epsilon) inv_var = nn.rsqrt(nn.add(variance, epsilon)) # batchnorm, I think? x = nn.sub(input, mean), x = nn.mul(x, inv_var) x = nn.mul(x, scale) x = nn.add(x, bias); return x I don't know why they didn't use nn.rsqrt, since it's a common operation. In fact it's so common it's a builtin op: https://www.tensorflow.org/api_docs/python/tf/raw_ops/Rsqrt https://www.tensorflow.org/api_docs/python/tf/raw_ops/Rsqrt I'm skeptical the reshaping is even necessary. Any reasonable framework should broadcast those shapes automatically.