| 198 | } |
| 199 | |
| 200 | void Compute(OpKernelContext* context) override { |
| 201 | TensorShape processing_shape, final_shape; |
| 202 | bool is_identity = true; |
| 203 | bool slice_dim0 = true; |
| 204 | bool is_simple_slice = true; |
| 205 | gtl::InlinedVector<int64, 4> begin; |
| 206 | gtl::InlinedVector<int64, 4> end; |
| 207 | gtl::InlinedVector<int64, 4> strides; |
| 208 | |
| 209 | TensorShape input_shape; |
| 210 | const Tensor& input_shape_tensor = context->input(0); |
| 211 | OP_REQUIRES( |
| 212 | context, input_shape_tensor.dims() == 1, |
| 213 | errors::InvalidArgument("shape must be 1-D, got shape.shape = ", |
| 214 | input_shape_tensor.shape().DebugString())); |
| 215 | if (input_shape_tensor.dtype() == DT_INT32) { |
| 216 | OP_REQUIRES_OK( |
| 217 | context, TensorShapeUtils::MakeShape(input_shape_tensor.vec<int32>(), |
| 218 | &input_shape)); |
| 219 | } else if (input_shape_tensor.dtype() == DT_INT64) { |
| 220 | OP_REQUIRES_OK( |
| 221 | context, TensorShapeUtils::MakeShape(input_shape_tensor.vec<int64>(), |
| 222 | &input_shape)); |
| 223 | } else { |
| 224 | LOG(FATAL) << "shape must have type int32 or int64."; |
| 225 | } |
| 226 | |
| 227 | OP_REQUIRES_OK( |
| 228 | context, |
| 229 | ValidateStridedSliceOp( |
| 230 | &context->input(1), &context->input(2), context->input(3), |
| 231 | input_shape, begin_mask, end_mask, ellipsis_mask, new_axis_mask, |
| 232 | shrink_axis_mask, &processing_shape, &final_shape, &is_identity, |
| 233 | &is_simple_slice, &slice_dim0, &begin, &end, &strides)); |
| 234 | |
| 235 | // Check to make sure dy is consistent with the original slice |
| 236 | TensorShape dy_shape = context->input(4).shape(); |
| 237 | OP_REQUIRES( |
| 238 | context, final_shape == dy_shape, |
| 239 | errors::InvalidArgument("shape of dy was ", dy_shape.DebugString(), |
| 240 | " instead of ", final_shape.DebugString())); |
| 241 | |
| 242 | if (!context->status().ok()) return; |
| 243 | |
| 244 | // const int input_dims = input.dims(); |
| 245 | const int processing_dims = processing_shape.dims(); |
| 246 | Tensor* result = nullptr; |
| 247 | OP_REQUIRES_OK(context, context->allocate_output(0, input_shape, &result)); |
| 248 | |
| 249 | if (processing_shape.dims() == 0) { |
| 250 | auto in = context->input(4); |
| 251 | OP_REQUIRES(context, result->CopyFrom(in, processing_shape), |
| 252 | errors::Internal("Copy failed")); |
| 253 | return; |
| 254 | } |
| 255 | |
| 256 | #define HANDLE_DIM(NDIM) \ |
| 257 | if (processing_dims == NDIM) { \ |
nothing calls this directly
no test coverage detected