| 2047 | } |
| 2048 | |
| 2049 | XlaOp XlaBuilder::BatchNormTraining(XlaOp operand, XlaOp scale, XlaOp offset, |
| 2050 | XlaOp side_input, float epsilon, |
| 2051 | int64 feature_index, |
| 2052 | size_t reserve_space_size, |
| 2053 | bool use_reserve_space, |
| 2054 | bool is_activation_relu) { |
| 2055 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 2056 | HloInstructionProto instr; |
| 2057 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 2058 | TF_ASSIGN_OR_RETURN(const Shape* scale_shape, GetShapePtr(scale)); |
| 2059 | TF_ASSIGN_OR_RETURN(const Shape* offset_shape, GetShapePtr(offset)); |
| 2060 | TF_ASSIGN_OR_RETURN( |
| 2061 | Shape shape, ShapeInference::InferBatchNormTrainingShape( |
| 2062 | *operand_shape, *scale_shape, *offset_shape, |
| 2063 | feature_index, reserve_space_size, use_reserve_space)); |
| 2064 | *instr.mutable_shape() = shape.ToProto(); |
| 2065 | |
| 2066 | instr.set_epsilon(epsilon); |
| 2067 | instr.set_feature_index(feature_index); |
| 2068 | instr.set_is_activation_relu(is_activation_relu); |
| 2069 | |
| 2070 | std::vector<XlaOp> operands = {operand, scale, offset}; |
| 2071 | if (!side_input.IsUninitialized()) { |
| 2072 | operands.push_back(side_input); |
| 2073 | } |
| 2074 | return AddInstruction(std::move(instr), HloOpcode::kBatchNormTraining, |
| 2075 | operands); |
| 2076 | }); |
| 2077 | } |
| 2078 | |
| 2079 | XlaOp XlaBuilder::BatchNormInference(XlaOp operand, XlaOp scale, XlaOp offset, |
| 2080 | XlaOp mean, XlaOp variance, float epsilon, |
no test coverage detected