| 3224 | namespace { |
| 3225 | |
| 3226 | Status ScatterNdShapeHelper(InferenceContext* c, ShapeHandle indices_shape, |
| 3227 | ShapeHandle updates_shape, |
| 3228 | ShapeHandle output_shape) { |
| 3229 | if (c->Value(c->NumElements(output_shape)) == 0 && |
| 3230 | (c->Value(c->NumElements(indices_shape)) > 0 || |
| 3231 | c->Value(c->NumElements(updates_shape)) > 0)) { |
| 3232 | return errors::InvalidArgument( |
| 3233 | "Indices and updates specified for empty output shape"); |
| 3234 | } |
| 3235 | |
| 3236 | if (c->RankKnown(indices_shape) && c->RankKnown(updates_shape)) { |
| 3237 | const int64 outer_dims = c->Rank(indices_shape) - 1; |
| 3238 | const DimensionHandle ixdim = c->Dim(indices_shape, -1); |
| 3239 | |
| 3240 | // We can only do more validation if the last dimension of indices |
| 3241 | // is a known value. |
| 3242 | if (c->ValueKnown(ixdim)) { |
| 3243 | int64 ix = c->Value(ixdim); |
| 3244 | ShapeHandle unused; |
| 3245 | ShapeHandle prefix_indices; |
| 3246 | TF_RETURN_IF_ERROR( |
| 3247 | c->Subshape(indices_shape, 0, outer_dims, &prefix_indices)); |
| 3248 | ShapeHandle prefix_updates; |
| 3249 | TF_RETURN_IF_ERROR( |
| 3250 | c->Subshape(updates_shape, 0, outer_dims, &prefix_updates)); |
| 3251 | |
| 3252 | Status s = c->Merge(prefix_indices, prefix_updates, &unused); |
| 3253 | if (!s.ok()) { |
| 3254 | return errors::InvalidArgument( |
| 3255 | "The outer ", outer_dims, |
| 3256 | " dimensions of indices.shape=", c->DebugString(indices_shape), |
| 3257 | " must match the outer ", outer_dims, |
| 3258 | " dimensions of updates.shape=", c->DebugString(updates_shape), |
| 3259 | ": ", s.error_message()); |
| 3260 | } |
| 3261 | |
| 3262 | ShapeHandle suffix_output; |
| 3263 | TF_RETURN_IF_ERROR(c->Subshape(output_shape, ix, &suffix_output)); |
| 3264 | ShapeHandle suffix_updates; |
| 3265 | TF_RETURN_IF_ERROR( |
| 3266 | c->Subshape(updates_shape, outer_dims, &suffix_updates)); |
| 3267 | s = c->Merge(suffix_output, suffix_updates, &unused); |
| 3268 | if (!s.ok()) { |
| 3269 | return errors::InvalidArgument( |
| 3270 | "The inner ", c->Rank(output_shape) - ix, |
| 3271 | " dimensions of output.shape=", c->DebugString(output_shape), |
| 3272 | " must match the inner ", c->Rank(updates_shape) - outer_dims, |
| 3273 | " dimensions of updates.shape=", c->DebugString(updates_shape), |
| 3274 | ": ", s.error_message()); |
| 3275 | } |
| 3276 | } |
| 3277 | } |
| 3278 | |
| 3279 | c->set_output(0, output_shape); |
| 3280 | return Status::OK(); |
| 3281 | } |
| 3282 | |
| 3283 | Status ScatterNdShape(InferenceContext* c) { |
no test coverage detected