Searched refs:activation_shape (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | batchnorm_expander.cc | 462 const Shape activation_shape = activation->shape(); in HandleBatchNormGrad() local
463 PrimitiveType ptype = activation_shape.element_type(); in HandleBatchNormGrad()
484 HloInstruction::CreateBroadcast(activation_shape, epsilon_scalar, {})); in HandleBatchNormGrad()
490 for (int64 i = 0; i < activation_shape.rank(); ++i) { in HandleBatchNormGrad()
497 activation_shape, scale, {feature_index})); in HandleBatchNormGrad()
499 activation_shape, variance, {feature_index})); in HandleBatchNormGrad()
503 HloInstruction::CreateBroadcast(activation_shape, mean, {feature_index})); in HandleBatchNormGrad()
507 add(Rsqrt(add_binary(activation_shape, HloOpcode::kAdd, in HandleBatchNormGrad()
517 activation_shape, HloOpcode::kSubtract, activation, mean_broadcasted); in HandleBatchNormGrad()
521 add_binary(activation_shape, HloOpcode::kMultiply, grad_output, in HandleBatchNormGrad()
[all …]
|