| 273 | } |
| 274 | |
| 275 | xla::ProgramShape XlaCompiledProgramShape( |
| 276 | const xla::XlaComputation& computation, |
| 277 | const xla::ProgramShape& input_program_shape) { |
| 278 | se::Platform* platform = |
| 279 | xla::PlatformUtil::GetPlatform(*xla_platform_ptr).ValueOrDie(); |
| 280 | xla::LocalClient* client = |
| 281 | xla::ClientLibrary::GetOrCreateLocalClient(platform).ValueOrDie(); |
| 282 | xla::ExecutableBuildOptions exec_options; |
| 283 | exec_options.set_result_layout(input_program_shape.result()); |
| 284 | std::vector<const xla::Shape*> parameters_shapes; |
| 285 | for (int64 i = 0; i < input_program_shape.parameters_size(); ++i) { |
| 286 | parameters_shapes.push_back(&input_program_shape.parameters(i)); |
| 287 | } |
| 288 | std::vector<std::unique_ptr<xla::LocalExecutable>> local_executables = |
| 289 | client->Compile(computation, parameters_shapes, exec_options) |
| 290 | .ConsumeValueOrDie(); |
| 291 | EXPECT_EQ(local_executables.size(), 1); |
| 292 | std::unique_ptr<xla::LocalExecutable> local_executable = |
| 293 | std::move(local_executables[0]); |
| 294 | return local_executable->executable() |
| 295 | ->module() |
| 296 | .entry_computation() |
| 297 | ->ComputeProgramShape(); |
| 298 | } |
| 299 | |
| 300 | TEST(RawApiTest, AllocFromTensor) { |
| 301 | xla::Literal literal = |
no test coverage detected