| 198 | } |
| 199 | |
| 200 | Status XlaResource::SetFromPack(const std::set<string>& gradient_sources, |
| 201 | const xla::XlaOp& pack, |
| 202 | xla::XlaBuilder* builder) { |
| 203 | if (gradient_sources.empty()) { |
| 204 | if (!initialized()) { |
| 205 | initial_value_ = pack; |
| 206 | } |
| 207 | value_ = pack; |
| 208 | } else { |
| 209 | TF_RET_CHECK(kind_ == kTensorArray); |
| 210 | int pos = 0; |
| 211 | auto v = xla::GetTupleElement(pack, pos++); |
| 212 | if (!initialized()) { |
| 213 | initial_value_ = v; |
| 214 | } |
| 215 | value_ = v; |
| 216 | |
| 217 | for (const auto& source : gradient_sources) { |
| 218 | XlaResource* gradient; |
| 219 | TF_RETURN_IF_ERROR( |
| 220 | GetOrCreateTensorArrayGradient(source, builder, &gradient)); |
| 221 | auto v = xla::GetTupleElement(pack, pos++); |
| 222 | if (!gradient->initialized()) { |
| 223 | gradient->initial_value_ = v; |
| 224 | } |
| 225 | gradient->value_ = v; |
| 226 | } |
| 227 | } |
| 228 | return Status::OK(); |
| 229 | } |
| 230 | |
| 231 | } // namespace tensorflow |
no test coverage detected