| 2077 | } |
| 2078 | |
| 2079 | XlaOp XlaBuilder::BatchNormInference(XlaOp operand, XlaOp scale, XlaOp offset, |
| 2080 | XlaOp mean, XlaOp variance, float epsilon, |
| 2081 | int64 feature_index) { |
| 2082 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 2083 | HloInstructionProto instr; |
| 2084 | |
| 2085 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 2086 | TF_ASSIGN_OR_RETURN(const Shape* scale_shape, GetShapePtr(scale)); |
| 2087 | TF_ASSIGN_OR_RETURN(const Shape* offset_shape, GetShapePtr(offset)); |
| 2088 | TF_ASSIGN_OR_RETURN(const Shape* mean_shape, GetShapePtr(mean)); |
| 2089 | TF_ASSIGN_OR_RETURN(const Shape* variance_shape, GetShapePtr(variance)); |
| 2090 | TF_ASSIGN_OR_RETURN(Shape shape, |
| 2091 | ShapeInference::InferBatchNormInferenceShape( |
| 2092 | *operand_shape, *scale_shape, *offset_shape, |
| 2093 | *mean_shape, *variance_shape, feature_index)); |
| 2094 | *instr.mutable_shape() = shape.ToProto(); |
| 2095 | |
| 2096 | instr.set_epsilon(epsilon); |
| 2097 | instr.set_feature_index(feature_index); |
| 2098 | |
| 2099 | return AddInstruction(std::move(instr), HloOpcode::kBatchNormInference, |
| 2100 | {operand, scale, offset, mean, variance}); |
| 2101 | }); |
| 2102 | } |
| 2103 | |
| 2104 | XlaOp XlaBuilder::BatchNormGrad(XlaOp operand, XlaOp scale, XlaOp batch_mean, |
| 2105 | XlaOp batch_var, XlaOp grad_output, |
no test coverage detected