| 183 | } |
| 184 | |
| 185 | Status XlaResource::Pack(xla::XlaOp* pack, xla::XlaBuilder* builder) const { |
| 186 | if (tensor_array_gradients_.empty()) { |
| 187 | *pack = value_; |
| 188 | } else { |
| 189 | TF_RET_CHECK(kind_ == kTensorArray); |
| 190 | std::vector<xla::XlaOp> elems; |
| 191 | elems.push_back(value_); |
| 192 | for (const auto& gradient : tensor_array_gradients_) { |
| 193 | elems.push_back(gradient.second->value_); |
| 194 | } |
| 195 | *pack = xla::Tuple(builder, elems); |
| 196 | } |
| 197 | return Status::OK(); |
| 198 | } |
| 199 | |
| 200 | Status XlaResource::SetFromPack(const std::set<string>& gradient_sources, |
| 201 | const xla::XlaOp& pack, |