| 1374 | } |
| 1375 | |
| 1376 | Status XlaCompiler::GetDeviceToHostShapes( |
| 1377 | const string& key, std::vector<TensorShape>* shapes) const { |
| 1378 | const auto iter = host_compute_sends_.find(key); |
| 1379 | if (iter == host_compute_sends_.end()) { |
| 1380 | return errors::InvalidArgument( |
| 1381 | "No host compute send shapes registered for key ", key); |
| 1382 | } |
| 1383 | shapes->clear(); |
| 1384 | for (int i = 0; i < iter->second.metadata_size(); ++i) { |
| 1385 | TensorShape shape(iter->second.metadata(i).shape()); |
| 1386 | shapes->push_back(shape); |
| 1387 | } |
| 1388 | return Status::OK(); |
| 1389 | } |
| 1390 | |
| 1391 | Status XlaCompiler::SetHostToDeviceMetadata( |
| 1392 | const string& key, absl::Span<const DataType> types, |