def compute_expect(data_sum, in_bn, head_active, in_active, layout): out_dtype = data_sum.dtype relugrad = relu_grad_np(head_active, in_active).astype(out_dtype) inbn_cast = in_bn.astype(out_dtype) bn_beta_ad = bn_beta_grad_np(relugrad, layout) bn_gamma_ad = bn_gamma_grad_np(relugrad, inbn_cast, data_sum, layout) return [bn_gamma_ad, bn_beta_ad]
def bn_update_grad(head, data_sum, in_bn, layout): out_dtype = data_sum.dtype head_cast = head.astype(out_dtype) inbn_cast = in_bn.astype(out_dtype) bn_beta_ad = bn_beta_grad_np(head_cast, layout) bn_gamma_ad = bn_gamma_grad_np(head_cast, inbn_cast, data_sum, layout) return [bn_beta_ad, bn_gamma_ad]