| 120 | } |
| 121 | |
| 122 | void Compute(OpKernelContext* context) override { |
| 123 | const int kFixedInputs = 4; // Prefix, tensor names, is_sparse |
| 124 | const Tensor& prefix = context->input(0); |
| 125 | const string& prefix_string = prefix.scalar<string>()(); |
| 126 | const Tensor& tensor_names = context->input(1); |
| 127 | const Tensor& shape_and_slices = context->input(2); |
| 128 | const Tensor& is_sparse = context->input(3); |
| 129 | const int num_tensors = static_cast<int>(tensor_names.NumElements()); |
| 130 | const auto& tensor_names_flat = tensor_names.flat<string>(); |
| 131 | const auto& is_sparse_flat = is_sparse.flat<bool>(); |
| 132 | const auto& shape_and_slices_flat = shape_and_slices.flat<string>(); |
| 133 | LOG(INFO) << "prefix_string: " << prefix_string |
| 134 | << "num tensors:" << num_tensors; |
| 135 | auto rm = context->resource_manager(); |
| 136 | BundleWriter writer(Env::Default(), prefix_string); |
| 137 | |
| 138 | for (int i = 0; i < num_tensors; i++) { |
| 139 | const string& tensor_name = tensor_names_flat(i); |
| 140 | if (is_sparse_flat(i)) { |
| 141 | IndicesIncrRecorder<int64>* sparse_incr_res = nullptr; |
| 142 | rm->Lookup("", tensor_name + "_sparse_incr", &sparse_incr_res); |
| 143 | if (sparse_incr_res != nullptr) { |
| 144 | DumpIncrSparse<int64>(context, i, kFixedInputs, |
| 145 | tensor_name, &writer, sparse_incr_res); |
| 146 | } else { |
| 147 | IndicesIncrRecorder<int32>* sparse_incr_res = nullptr; |
| 148 | rm->Lookup("", tensor_name + "_sparse_incr", |
| 149 | &sparse_incr_res); |
| 150 | if (sparse_incr_res != nullptr) { |
| 151 | DumpIncrSparse<int32>(context, i, kFixedInputs, |
| 152 | tensor_name, &writer, sparse_incr_res); |
| 153 | } else { |
| 154 | LOG(WARNING) << tensor_name << "_sparse_incr" |
| 155 | << " Resource NOT FOUND"; |
| 156 | } |
| 157 | } |
| 158 | } else { |
| 159 | const Tensor& tensor = context->input(i + kFixedInputs); |
| 160 | |
| 161 | if (!shape_and_slices_flat(i).empty()) { |
| 162 | const string& shape_spec = shape_and_slices_flat(i); |
| 163 | TensorShape shape; |
| 164 | TensorSlice slice(tensor.dims()); |
| 165 | TensorShape slice_shape; |
| 166 | |
| 167 | OP_REQUIRES_OK(context, checkpoint::ParseShapeAndSlice( |
| 168 | shape_spec, &shape, &slice, &slice_shape)); |
| 169 | OP_REQUIRES(context, slice_shape.IsSameSize(tensor.shape()), |
| 170 | errors::InvalidArgument("Slice in shape_and_slice " |
| 171 | "specification does not match the " |
| 172 | "shape of the tensor to save: ", |
| 173 | shape_spec, ", tensor: ", |
| 174 | tensor.shape().DebugString())); |
| 175 | |
| 176 | OP_REQUIRES_OK(context, |
| 177 | writer.AddSlice(tensor_name, shape, slice, tensor)); |
| 178 | } else { |
| 179 | OP_REQUIRES_OK(context, writer.Add(tensor_name, tensor)); |
nothing calls this directly
no test coverage detected