| 363 | } |
| 364 | |
| 365 | StatusOr<HloInstruction*> MakeSelectHlo(HloInstruction* pred, |
| 366 | HloInstruction* on_true, |
| 367 | HloInstruction* on_false, |
| 368 | HloInstruction* derived_from) { |
| 369 | HloComputation* computation = pred->parent(); |
| 370 | DCHECK_EQ(computation, on_true->parent()); |
| 371 | DCHECK_EQ(computation, on_false->parent()); |
| 372 | Shape op_shape = on_true->shape(); |
| 373 | if (ShapeUtil::IsScalar(pred->shape())) { |
| 374 | if (!ShapeUtil::IsScalar(op_shape) && !op_shape.IsTuple()) { |
| 375 | // If the output is not scalar, we need to broadcast the condition |
| 376 | // to match the contract of kSelect. For tuples, we use kTupleSelect |
| 377 | // which expects the condition to be a scalar. |
| 378 | pred = computation->AddInstruction(HloInstruction::CreateBroadcast( |
| 379 | ShapeUtil::ChangeElementType(op_shape, PrimitiveType::PRED), pred, |
| 380 | {})); |
| 381 | if (derived_from) { |
| 382 | derived_from->SetupDerivedInstruction(pred); |
| 383 | } |
| 384 | } |
| 385 | } |
| 386 | HloOpcode select_op_code = |
| 387 | op_shape.IsTuple() ? HloOpcode::kTupleSelect : HloOpcode::kSelect; |
| 388 | TF_ASSIGN_OR_RETURN(Shape select_shape, |
| 389 | ShapeInference::InferTernaryOpShape(select_op_code, pred, |
| 390 | on_true, on_false)); |
| 391 | HloInstruction* select = |
| 392 | computation->AddInstruction(HloInstruction::CreateTernary( |
| 393 | select_shape, select_op_code, pred, on_true, on_false)); |
| 394 | if (derived_from) { |
| 395 | derived_from->SetupDerivedInstruction(select); |
| 396 | } |
| 397 | return select; |
| 398 | } |
| 399 | |
| 400 | StatusOr<HloInstruction*> MakeSortHlo( |
| 401 | const Shape& sort_shape, absl::Span<HloInstruction* const> operands, |
no test coverage detected