| 243 | }); |
| 244 | |
| 245 | Status TensorListConcatShapeInference( |
| 246 | shape_inference::InferenceContext* c, |
| 247 | shape_inference::ShapeHandle element_shape) { |
| 248 | DataType element_dtype; |
| 249 | TF_RETURN_IF_ERROR(c->GetAttr("element_dtype", &element_dtype)); |
| 250 | auto* handle_data = c->input_handle_shapes_and_types(0); |
| 251 | if (handle_data != nullptr && handle_data->size() > 1) { |
| 252 | return errors::InvalidArgument( |
| 253 | "Trying to read from list with wrong variant data."); |
| 254 | } |
| 255 | if (handle_data != nullptr && handle_data->size() == 1) { |
| 256 | const shape_inference::ShapeAndType& list_shape_type = (*handle_data)[0]; |
| 257 | if (list_shape_type.dtype != element_dtype) { |
| 258 | return errors::InvalidArgument( |
| 259 | "Trying to read from list with wrong element dtype. List has " |
| 260 | "type ", |
| 261 | DataTypeString(list_shape_type.dtype), " but expected type ", |
| 262 | DataTypeString(element_dtype)); |
| 263 | } |
| 264 | shape_inference::ShapeHandle merged; |
| 265 | TF_RETURN_IF_ERROR(c->Merge(element_shape, list_shape_type.shape, &merged)); |
| 266 | element_shape = merged; |
| 267 | } |
| 268 | if (c->RankKnown(element_shape)) { |
| 269 | shape_inference::ShapeHandle result; |
| 270 | TF_RETURN_IF_ERROR(c->Subshape(element_shape, 1, &result)); |
| 271 | TF_RETURN_IF_ERROR( |
| 272 | c->Concatenate(c->MakeShape({c->UnknownDim()}), result, &result)); |
| 273 | c->set_output(0, result); |
| 274 | } else { |
| 275 | c->set_output(0, c->UnknownShape()); |
| 276 | } |
| 277 | c->set_output(1, c->MakeShape({c->UnknownDim()})); |
| 278 | return Status::OK(); |
| 279 | } |
| 280 | |
| 281 | REGISTER_OP("TensorListConcat") |
| 282 | .Input("input_handle: variant") |
no test coverage detected