| 2102 | } |
| 2103 | |
| 2104 | XlaOp XlaBuilder::BatchNormGrad(XlaOp operand, XlaOp scale, XlaOp batch_mean, |
| 2105 | XlaOp batch_var, XlaOp grad_output, |
| 2106 | XlaOp reserve_space, float epsilon, |
| 2107 | int64 feature_index) { |
| 2108 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 2109 | HloInstructionProto instr; |
| 2110 | |
| 2111 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 2112 | TF_ASSIGN_OR_RETURN(const Shape* scale_shape, GetShapePtr(scale)); |
| 2113 | TF_ASSIGN_OR_RETURN(const Shape* batch_mean_shape, GetShapePtr(batch_mean)); |
| 2114 | TF_ASSIGN_OR_RETURN(const Shape* batch_var_shape, GetShapePtr(batch_var)); |
| 2115 | TF_ASSIGN_OR_RETURN(const Shape* grad_output_shape, |
| 2116 | GetShapePtr(grad_output)); |
| 2117 | TF_ASSIGN_OR_RETURN( |
| 2118 | Shape shape, ShapeInference::InferBatchNormGradShape( |
| 2119 | *operand_shape, *scale_shape, *batch_mean_shape, |
| 2120 | *batch_var_shape, *grad_output_shape, feature_index)); |
| 2121 | *instr.mutable_shape() = shape.ToProto(); |
| 2122 | |
| 2123 | instr.set_epsilon(epsilon); |
| 2124 | instr.set_feature_index(feature_index); |
| 2125 | std::vector<XlaOp> operands = {operand, scale, batch_mean, batch_var, |
| 2126 | grad_output}; |
| 2127 | if (!reserve_space.IsUninitialized()) { |
| 2128 | operands.push_back(reserve_space); |
| 2129 | } |
| 2130 | return AddInstruction(std::move(instr), HloOpcode::kBatchNormGrad, |
| 2131 | operands); |
| 2132 | }); |
| 2133 | } |
| 2134 | |
| 2135 | XlaOp XlaBuilder::CrossReplicaSum( |
| 2136 | XlaOp operand, absl::Span<const ReplicaGroup> replica_groups) { |
no test coverage detected