| 689 | } |
| 690 | |
| 691 | XlaOp XlaBuilder::Call(const XlaComputation& computation, |
| 692 | absl::Span<const XlaOp> operands) { |
| 693 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 694 | HloInstructionProto instr; |
| 695 | std::vector<const Shape*> operand_shape_ptrs; |
| 696 | TF_ASSIGN_OR_RETURN(const auto& operand_shapes, GetOperandShapes(operands)); |
| 697 | absl::c_transform(operand_shapes, std::back_inserter(operand_shape_ptrs), |
| 698 | [](const Shape& shape) { return &shape; }); |
| 699 | TF_ASSIGN_OR_RETURN(const ProgramShape& called_program_shape, |
| 700 | computation.GetProgramShape()); |
| 701 | TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferCallShape( |
| 702 | operand_shape_ptrs, |
| 703 | /*to_apply=*/called_program_shape)); |
| 704 | *instr.mutable_shape() = shape.ToProto(); |
| 705 | |
| 706 | AddCalledComputation(computation, &instr); |
| 707 | |
| 708 | return AddInstruction(std::move(instr), HloOpcode::kCall, operands); |
| 709 | }); |
| 710 | } |
| 711 | |
| 712 | XlaOp XlaBuilder::Parameter( |
| 713 | int64 parameter_number, const Shape& shape, const string& name, |
no test coverage detected