| 37 | explicit SliceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} |
| 38 | |
| 39 | void Compile(XlaOpKernelContext* ctx) override { |
| 40 | const TensorShape input_shape = ctx->InputShape(0); |
| 41 | const TensorShape begin_tensor_shape = ctx->InputShape(1); |
| 42 | const TensorShape size_tensor_shape = ctx->InputShape(2); |
| 43 | |
| 44 | const int input_dims = input_shape.dims(); |
| 45 | OP_REQUIRES( |
| 46 | ctx, |
| 47 | TensorShapeUtils::IsVector(begin_tensor_shape) && |
| 48 | TensorShapeUtils::IsVector(size_tensor_shape) && |
| 49 | begin_tensor_shape.num_elements() == input_dims && |
| 50 | size_tensor_shape.num_elements() == input_dims, |
| 51 | errors::InvalidArgument( |
| 52 | "Expected begin and size arguments to be 1-D tensors of size ", |
| 53 | input_dims, ", but got shapes ", begin_tensor_shape.DebugString(), |
| 54 | " and ", size_tensor_shape.DebugString(), " instead.")); |
| 55 | |
| 56 | std::vector<int64> begin; |
| 57 | std::vector<int64> size; |
| 58 | OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(2, &size)); |
| 59 | if (ctx->ConstantInputAsIntVector(1, &begin).ok()) { |
| 60 | // `begin` is a compile-time constant. |
| 61 | for (int i = 0; i < input_dims; ++i) { |
| 62 | if (size[i] == -1) { |
| 63 | // A size[i] of -1 means "all elements from begin[i] to dim_size(i)". |
| 64 | size[i] = input_shape.dim_size(i) - begin[i]; |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | for (int i = 0; i < input_dims; ++i) { |
| 69 | int64 b = begin[i]; |
| 70 | int64 s = size[i]; |
| 71 | if (input_shape.dim_size(i) == 0) { |
| 72 | OP_REQUIRES(ctx, b == 0 && s == 0, |
| 73 | errors::InvalidArgument( |
| 74 | "Expected begin[", i, "] == 0 (got ", b, |
| 75 | ") and size[", i, "] == 0 ", "(got ", s, ") when ", |
| 76 | "input_shape.dim_size(", i, ") == 0")); |
| 77 | } else { |
| 78 | OP_REQUIRES(ctx, 0 <= b && b <= input_shape.dim_size(i), |
| 79 | errors::InvalidArgument("Expected begin[", i, "] in [0, ", |
| 80 | input_shape.dim_size(i), |
| 81 | "], but got ", b)); |
| 82 | OP_REQUIRES(ctx, 0 <= s && b + s <= input_shape.dim_size(i), |
| 83 | errors::InvalidArgument("Expected size[", i, "] in [0, ", |
| 84 | input_shape.dim_size(i) - b, |
| 85 | "], but ", "got ", s)); |
| 86 | } |
| 87 | } |
| 88 | |
| 89 | std::vector<int64> limits; |
| 90 | limits.reserve(begin.size()); |
| 91 | for (int i = 0; i < begin.size(); ++i) { |
| 92 | limits.push_back(begin[i] + size[i]); |
| 93 | } |
| 94 | std::vector<int64> strides(begin.size(), 1); |
| 95 | ctx->SetOutput(0, xla::Slice(ctx->Input(0), begin, limits, strides)); |
| 96 | } else { |
nothing calls this directly
no test coverage detected