| 166 | .SetShapeFn(shape_inference::NoOutputs); |
| 167 | |
| 168 | Status CreateAssignShapeFn(InferenceContext* c) { |
| 169 | std::vector<ShapeAndType> handle_shape_and_type; |
| 170 | TF_RETURN_IF_ERROR(shape_inference::ValidateVariableResourceHandle( |
| 171 | c, &handle_shape_and_type)); |
| 172 | |
| 173 | ShapeHandle value_shape = c->input(1); |
| 174 | ShapeHandle unused; |
| 175 | TF_RETURN_IF_ERROR( |
| 176 | c->Merge(handle_shape_and_type[0].shape, value_shape, &unused)); |
| 177 | |
| 178 | if (handle_shape_and_type[0].dtype == DT_VARIANT && |
| 179 | handle_shape_and_type.size() > 1 && |
| 180 | c->input_handle_shapes_and_types(1) != nullptr) { |
| 181 | auto* value_handle_shape_and_type = c->input_handle_shapes_and_types(1); |
| 182 | if (value_handle_shape_and_type->size() != |
| 183 | handle_shape_and_type.size() - 1) { |
| 184 | return errors::InvalidArgument( |
| 185 | "Incompatible handle variant shape_and_type size and input " |
| 186 | "shape_and_type size: ", |
| 187 | handle_shape_and_type.size() - 1, " vs. ", |
| 188 | value_handle_shape_and_type->size()); |
| 189 | } |
| 190 | } |
| 191 | return Status::OK(); |
| 192 | } |
| 193 | |
| 194 | REGISTER_OP("AssignVariableOp") |
| 195 | .Input("resource: resource") |
nothing calls this directly
no test coverage detected