Lines Matching defs:feature_index
165 int64 feature_index = batch_norm->feature_index();
166 const int64 feature_count = operand_shape.dimensions(feature_index);
190 if (i != feature_index) {
196 HloInstruction::CreateBroadcast(operand_shape, scale, {feature_index}));
199 HloInstruction::CreateBroadcast(operand_shape, offset, {feature_index}));
236 HloInstruction::CreateBroadcast(operand_shape, mean, {feature_index}));
251 add(HloInstruction::CreateBroadcast(operand_shape, var, {feature_index}));
312 int64 feature_index = batch_norm->feature_index();
329 if (i != feature_index) {
343 HloInstruction::CreateBroadcast(operand_shape, scale, {feature_index}));
346 HloInstruction::CreateBroadcast(operand_shape, offset, {feature_index}));
349 HloInstruction::CreateBroadcast(operand_shape, mean, {feature_index}));
352 add(HloInstruction::CreateBroadcast(operand_shape, var, {feature_index}));
436 int64 feature_index = batch_norm->feature_index();
439 const int64 feature_count = activation_shape.dimensions(feature_index);
464 if (i != feature_index) {
470 activation_shape, scale, {feature_index}));
472 activation_shape, variance, {feature_index}));
476 HloInstruction::CreateBroadcast(activation_shape, mean, {feature_index}));
536 {feature_index}));
541 {feature_index}));
