| 2009 | } |
| 2010 | |
| 2011 | Status ScatterNdUpdateShape(InferenceContext* c) { |
| 2012 | ShapeHandle input_shape = c->input(0); |
| 2013 | if (c->input_handle_shapes_and_types(0) != nullptr) { |
| 2014 | // This is called for tf.scatter_nd_update; input is a Variable handle. |
| 2015 | const auto& shape_and_type = *(c->input_handle_shapes_and_types(0)); |
| 2016 | if (shape_and_type.size() == 1) { |
| 2017 | input_shape = shape_and_type[0].shape; |
| 2018 | } |
| 2019 | } |
| 2020 | ShapeHandle indices_shape; |
| 2021 | TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(1), 1, &indices_shape)); |
| 2022 | ShapeHandle updates_shape; |
| 2023 | TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(2), 1, &updates_shape)); |
| 2024 | |
| 2025 | if (c->Value(c->NumElements(input_shape)) == 0 && |
| 2026 | (c->Value(c->NumElements(indices_shape)) > 0 || |
| 2027 | c->Value(c->NumElements(updates_shape)) > 0)) { |
| 2028 | return errors::InvalidArgument( |
| 2029 | "Indices and updates specified for empty output shape"); |
| 2030 | } |
| 2031 | |
| 2032 | if (c->RankKnown(indices_shape) && c->RankKnown(updates_shape)) { |
| 2033 | const int64 num_outer_dims = c->Rank(indices_shape) - 1; |
| 2034 | const DimensionHandle index_size = c->Dim(indices_shape, -1); |
| 2035 | |
| 2036 | // We can only do more validation if the last dimension of indices |
| 2037 | // is a known value. |
| 2038 | if (c->ValueKnown(index_size)) { |
| 2039 | const int64 ix = c->Value(index_size); |
| 2040 | ShapeHandle unused; |
| 2041 | ShapeHandle prefix_indices; |
| 2042 | TF_RETURN_IF_ERROR( |
| 2043 | c->Subshape(indices_shape, 0, num_outer_dims, &prefix_indices)); |
| 2044 | ShapeHandle prefix_updates; |
| 2045 | TF_RETURN_IF_ERROR( |
| 2046 | c->Subshape(updates_shape, 0, num_outer_dims, &prefix_updates)); |
| 2047 | |
| 2048 | Status s = c->Merge(prefix_indices, prefix_updates, &unused); |
| 2049 | if (!s.ok()) { |
| 2050 | return errors::InvalidArgument( |
| 2051 | "The outer ", num_outer_dims, |
| 2052 | " dimensions of indices.shape=", c->DebugString(indices_shape), |
| 2053 | " must match the outer ", num_outer_dims, |
| 2054 | " dimensions of updates.shape=", c->DebugString(updates_shape), |
| 2055 | ": ", s.error_message()); |
| 2056 | } |
| 2057 | |
| 2058 | ShapeHandle input_suffix; |
| 2059 | TF_RETURN_IF_ERROR(c->Subshape(input_shape, ix, &input_suffix)); |
| 2060 | ShapeHandle suffix_updates; |
| 2061 | TF_RETURN_IF_ERROR( |
| 2062 | c->Subshape(updates_shape, num_outer_dims, &suffix_updates)); |
| 2063 | s = c->Merge(input_suffix, suffix_updates, &unused); |
| 2064 | if (!s.ok()) { |
| 2065 | return errors::InvalidArgument( |
| 2066 | "The inner ", c->Rank(input_shape) - ix, |
| 2067 | " dimensions of input.shape=", c->DebugString(input_shape), |
| 2068 | " must match the inner ", c->Rank(updates_shape) - num_outer_dims, |
nothing calls this directly
no test coverage detected