| 161 | } |
| 162 | |
| 163 | StatusOr<ProgramShape> XlaBuilder::GetProgramShape(int64 root_id) const { |
| 164 | TF_RETURN_IF_ERROR(first_error_); |
| 165 | TF_ASSIGN_OR_RETURN(const HloInstructionProto* root_proto, |
| 166 | LookUpInstructionByHandle(root_id)); |
| 167 | |
| 168 | ProgramShape program_shape; |
| 169 | |
| 170 | *program_shape.mutable_result() = Shape(root_proto->shape()); |
| 171 | |
| 172 | // Check that the parameter numbers are continuous from 0, and add parameter |
| 173 | // shapes and names to the program shape. |
| 174 | const int64 param_count = parameter_numbers_.size(); |
| 175 | for (int64 i = 0; i < param_count; i++) { |
| 176 | program_shape.add_parameters(); |
| 177 | program_shape.add_parameter_names(); |
| 178 | } |
| 179 | for (const HloInstructionProto& instr : instructions_) { |
| 180 | // Parameter number uniqueness is guaranteed in XlaBuilder::Parameter(). So |
| 181 | // to verify continuity, we just need to verify that every parameter is in |
| 182 | // the right range. |
| 183 | if (instr.opcode() == HloOpcodeString(HloOpcode::kParameter)) { |
| 184 | const int64 index = instr.parameter_number(); |
| 185 | TF_RET_CHECK(index >= 0 && index < param_count) |
| 186 | << "invalid parameter number: " << index; |
| 187 | *program_shape.mutable_parameters(index) = Shape(instr.shape()); |
| 188 | *program_shape.mutable_parameter_names(index) = instr.name(); |
| 189 | } |
| 190 | } |
| 191 | return program_shape; |
| 192 | } |
| 193 | |
| 194 | StatusOr<ProgramShape> XlaBuilder::GetProgramShape() const { |
| 195 | TF_RET_CHECK(!instructions_.empty()); |