| 2574 | } |
| 2575 | |
| 2576 | XlaOp XlaBuilder::GetDimensionSize(XlaOp operand, int64 dimension) { |
| 2577 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 2578 | HloInstructionProto instr; |
| 2579 | TF_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); |
| 2580 | TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferGetDimensionSizeShape( |
| 2581 | *operand_shape, dimension)); |
| 2582 | *instr.mutable_shape() = shape.ToProto(); |
| 2583 | instr.add_dimensions(dimension); |
| 2584 | return AddInstruction(std::move(instr), HloOpcode::kGetDimensionSize, |
| 2585 | {operand}); |
| 2586 | }); |
| 2587 | } |
| 2588 | |
| 2589 | XlaOp XlaBuilder::SetDimensionSize(XlaOp operand, XlaOp val, int64 dimension) { |
| 2590 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
no test coverage detected