| 1091 | } |
| 1092 | |
| 1093 | XlaOp XlaBuilder::GetTupleElement(XlaOp tuple_data, int64 index) { |
| 1094 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 1095 | HloInstructionProto instr; |
| 1096 | TF_ASSIGN_OR_RETURN(const Shape* tuple_shape, GetShapePtr(tuple_data)); |
| 1097 | if (!tuple_shape->IsTuple()) { |
| 1098 | return InvalidArgument( |
| 1099 | "Operand to GetTupleElement() is not a tuple; got %s", |
| 1100 | ShapeUtil::HumanString(*tuple_shape)); |
| 1101 | } |
| 1102 | if (index < 0 || index >= ShapeUtil::TupleElementCount(*tuple_shape)) { |
| 1103 | return InvalidArgument( |
| 1104 | "GetTupleElement() index (%d) out of range for tuple shape %s", index, |
| 1105 | ShapeUtil::HumanString(*tuple_shape)); |
| 1106 | } |
| 1107 | *instr.mutable_shape() = |
| 1108 | ShapeUtil::GetTupleElementShape(*tuple_shape, index).ToProto(); |
| 1109 | |
| 1110 | instr.set_tuple_index(index); |
| 1111 | |
| 1112 | return AddInstruction(std::move(instr), HloOpcode::kGetTupleElement, |
| 1113 | {tuple_data}); |
| 1114 | }); |
| 1115 | } |
| 1116 | |
| 1117 | XlaOp XlaBuilder::Dot(XlaOp lhs, XlaOp rhs, |
| 1118 | const PrecisionConfig* precision_config) { |