| 103 | } |
| 104 | |
| 105 | StatusOr<const Shape*> XlaBuilder::GetShapePtr(XlaOp op) const { |
| 106 | TF_RETURN_IF_ERROR(first_error_); |
| 107 | TF_RETURN_IF_ERROR(CheckOpBuilder(op)); |
| 108 | auto it = handle_to_index_.find(op.handle()); |
| 109 | if (it == handle_to_index_.end()) { |
| 110 | return InvalidArgument("No XlaOp with handle %d", op.handle()); |
| 111 | } |
| 112 | return instruction_shapes_.at(it->second).get(); |
| 113 | } |
| 114 | |
| 115 | StatusOr<Shape> XlaBuilder::GetShape(XlaOp op) const { |
| 116 | TF_ASSIGN_OR_RETURN(const Shape* shape, GetShapePtr(op)); |