| 160 | } |
| 161 | |
| 162 | Status SnapshotResourceVariables(OpKernelContext* ctx, |
| 163 | absl::Span<const int> variable_indices, |
| 164 | std::map<int, OptionalTensor>* result) { |
| 165 | std::vector<VariableInfo> variable_infos; |
| 166 | TF_RETURN_IF_ERROR( |
| 167 | GetVariableInfosFromCtxInputs(ctx, variable_indices, &variable_infos)); |
| 168 | TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(variable_infos))); |
| 169 | |
| 170 | for (int i = 0; i < variable_indices.size(); i++) { |
| 171 | if (variable_infos[i].var()) { |
| 172 | OptionalTensor& tensor = (*result)[variable_indices[i]]; |
| 173 | tensor.name = HandleFromInput(ctx, variable_indices[i]).name(); |
| 174 | tensor.present = true; |
| 175 | tensor.value = *variable_infos[i].var()->tensor(); |
| 176 | } else { |
| 177 | (*result)[variable_indices[i]] = OptionalTensor(); |
| 178 | } |
| 179 | } |
| 180 | return Status::OK(); |
| 181 | } |
| 182 | |
| 183 | XlaComputationLaunchContext::XlaComputationLaunchContext( |
| 184 | xla::LocalClient* client, se::DeviceMemoryAllocator* xla_allocator, |
no test coverage detected