| 88 | }); |
| 89 | |
| 90 | Status ValidateTableResourceHandle(InferenceContext* c, ShapeHandle keys, |
| 91 | const string& key_dtype_attr, |
| 92 | const string& value_dtype_attr, |
| 93 | bool is_lookup, |
| 94 | ShapeAndType* output_shape_and_type) { |
| 95 | auto* handle_data = c->input_handle_shapes_and_types(0); |
| 96 | if (handle_data == nullptr || handle_data->size() != 2) { |
| 97 | output_shape_and_type->shape = c->UnknownShape(); |
| 98 | output_shape_and_type->dtype = DT_INVALID; |
| 99 | } else { |
| 100 | const ShapeAndType& key_shape_and_type = (*handle_data)[0]; |
| 101 | const ShapeAndType& value_shape_and_type = (*handle_data)[1]; |
| 102 | DataType key_dtype; |
| 103 | TF_RETURN_IF_ERROR(c->GetAttr(key_dtype_attr, &key_dtype)); |
| 104 | if (key_shape_and_type.dtype != key_dtype) { |
| 105 | return errors::InvalidArgument( |
| 106 | "Trying to read value with wrong dtype. " |
| 107 | "Expected ", |
| 108 | DataTypeString(key_shape_and_type.dtype), " got ", |
| 109 | DataTypeString(key_dtype)); |
| 110 | } |
| 111 | DataType value_dtype; |
| 112 | TF_RETURN_IF_ERROR(c->GetAttr(value_dtype_attr, &value_dtype)); |
| 113 | if (value_shape_and_type.dtype != value_dtype) { |
| 114 | return errors::InvalidArgument( |
| 115 | "Trying to read value with wrong dtype. " |
| 116 | "Expected ", |
| 117 | DataTypeString(value_shape_and_type.dtype), " got ", |
| 118 | DataTypeString(value_dtype)); |
| 119 | } |
| 120 | output_shape_and_type->dtype = value_shape_and_type.dtype; |
| 121 | |
| 122 | if (is_lookup) { |
| 123 | if (c->RankKnown(key_shape_and_type.shape) && c->RankKnown(keys)) { |
| 124 | int keys_rank = c->Rank(keys); |
| 125 | int key_suffix_rank = c->Rank(key_shape_and_type.shape); |
| 126 | if (keys_rank < key_suffix_rank) { |
| 127 | return errors::InvalidArgument( |
| 128 | "Expected keys to have suffix ", |
| 129 | c->DebugString(key_shape_and_type.shape), |
| 130 | " but saw shape: ", c->DebugString(keys)); |
| 131 | } |
| 132 | for (int d = 0; d < key_suffix_rank; d++) { |
| 133 | // Ensure the suffix of keys match what's in the Table. |
| 134 | DimensionHandle dim = c->Dim(key_shape_and_type.shape, d); |
| 135 | TF_RETURN_IF_ERROR( |
| 136 | c->ReplaceDim(keys, keys_rank - key_suffix_rank + d, dim, &keys)); |
| 137 | } |
| 138 | std::vector<DimensionHandle> keys_prefix_vec; |
| 139 | keys_prefix_vec.reserve(keys_rank - key_suffix_rank); |
| 140 | for (int d = 0; d < keys_rank - key_suffix_rank; ++d) { |
| 141 | keys_prefix_vec.push_back(c->Dim(keys, d)); |
| 142 | } |
| 143 | ShapeHandle keys_prefix = c->MakeShape(keys_prefix_vec); |
| 144 | TF_RETURN_IF_ERROR(c->Concatenate(keys_prefix, |
| 145 | value_shape_and_type.shape, |
| 146 | &output_shape_and_type->shape)); |
| 147 | } else { |
no test coverage detected