| 150 | } |
| 151 | |
| 152 | Status ValidateStridedSliceOp( |
| 153 | const Tensor* begin_tensor, const Tensor* end_tensor, |
| 154 | const Tensor& strides_tensor, const PartialTensorShape& input_shape, |
| 155 | int32 begin_mask_spec, int32 end_mask_spec, const int32 ellipsis_mask, |
| 156 | int32 new_axis_mask, int32 shrink_axis_mask, |
| 157 | PartialTensorShape* processing_shape, PartialTensorShape* final_shape, |
| 158 | bool* is_identity, bool* is_simple_slice, bool* slice_dim0, |
| 159 | gtl::InlinedVector<int64, 4>* begin, gtl::InlinedVector<int64, 4>* end, |
| 160 | gtl::InlinedVector<int64, 4>* strides) { |
| 161 | const bool begin_is_wrong = |
| 162 | begin_tensor != nullptr && |
| 163 | !(TensorShapeUtils::IsVector(begin_tensor->shape()) && |
| 164 | begin_tensor->NumElements() == strides_tensor.NumElements() && |
| 165 | begin_tensor->NumElements() < 32 /* using 32 bit masks */); |
| 166 | const bool end_is_wrong = |
| 167 | end_tensor != nullptr && |
| 168 | !(TensorShapeUtils::IsVector(end_tensor->shape()) && |
| 169 | end_tensor->NumElements() == strides_tensor.NumElements()); |
| 170 | if (begin_is_wrong || end_is_wrong || |
| 171 | !TensorShapeUtils::IsVector(strides_tensor.shape())) { |
| 172 | if (begin_tensor != nullptr && end_tensor != nullptr) { |
| 173 | return errors::InvalidArgument( |
| 174 | "Expected begin, end, and strides to be 1D equal size tensors, ", |
| 175 | "but got shapes ", begin_tensor->shape().DebugString(), ", ", |
| 176 | end_tensor->shape().DebugString(), ", and ", |
| 177 | strides_tensor.shape().DebugString(), " instead."); |
| 178 | } else { |
| 179 | return errors::InvalidArgument( |
| 180 | "Expected begin, end, and strides to be 1D equal size tensors, ", |
| 181 | "but got shape ", strides_tensor.shape().DebugString(), |
| 182 | " for strides."); |
| 183 | } |
| 184 | } |
| 185 | // Use bit compares to ensure ellipsis_mask is 0 or a power of 2 |
| 186 | // i.e. there exists only no more than one ellipsis |
| 187 | if (ellipsis_mask && ((ellipsis_mask & (ellipsis_mask - 1)) != 0)) { |
| 188 | return errors::InvalidArgument( |
| 189 | "Multiple ellipses in slice spec not allowed"); |
| 190 | } |
| 191 | |
| 192 | // Step 1: Account for ellipsis and new axis |
| 193 | // |
| 194 | // Check for ellipses and count how many non-newaxis' there are after |
| 195 | // TODO(aselle): Convert this to do a fast log2 followed by iteration |
| 196 | // counting ones in next guys |
| 197 | bool ellipsis_seen = false; |
| 198 | |
| 199 | StridedSliceSparseSpec sparse_spec = {strides_tensor.NumElements(), |
| 200 | 0, |
| 201 | begin_tensor, |
| 202 | end_tensor, |
| 203 | strides_tensor, |
| 204 | begin_mask_spec, |
| 205 | end_mask_spec, |
| 206 | ellipsis_mask, |
| 207 | new_axis_mask, |
| 208 | shrink_axis_mask}; |
| 209 | |