| 45 | } |
| 46 | |
| 47 | void Compute(OpKernelContext* context) override { |
| 48 | // Get inputs |
| 49 | const Tensor& input_tensor = context->input(0); |
| 50 | const Tensor& pos_tensor = context->input(1); |
| 51 | const Tensor& len_tensor = context->input(2); |
| 52 | const TensorShape& input_shape = input_tensor.shape(); |
| 53 | const TensorShape& pos_shape = pos_tensor.shape(); |
| 54 | |
| 55 | bool is_scalar = TensorShapeUtils::IsScalar(pos_shape); |
| 56 | |
| 57 | if (is_scalar || input_shape == pos_shape) { |
| 58 | // pos/len are either scalar or match the shape of input_tensor |
| 59 | // Do not need to do broadcasting |
| 60 | |
| 61 | // Reshape input |
| 62 | auto input = input_tensor.flat<tstring>(); |
| 63 | // Allocate output |
| 64 | Tensor* output_tensor = nullptr; |
| 65 | OP_REQUIRES_OK(context, |
| 66 | context->allocate_output("output", input_tensor.shape(), |
| 67 | &output_tensor)); |
| 68 | auto output = output_tensor->flat<tstring>(); |
| 69 | if (is_scalar) { |
| 70 | // Perform Op with scalar pos/len |
| 71 | const T pos = |
| 72 | tensorflow::internal::SubtleMustCopy(pos_tensor.scalar<T>()()); |
| 73 | const T len = |
| 74 | tensorflow::internal::SubtleMustCopy(len_tensor.scalar<T>()()); |
| 75 | for (size_t i = 0; i < input_tensor.NumElements(); ++i) { |
| 76 | StringPiece in(input(i)); |
| 77 | T byte_pos = pos; |
| 78 | T byte_len = len; |
| 79 | switch (unit_) { |
| 80 | case CharUnit::UTF8_CHAR: |
| 81 | OP_REQUIRES( |
| 82 | context, UpdatePosAndLenForUtf8(in, &byte_pos, &byte_len), |
| 83 | errors::InvalidArgument("pos ", pos, " out of range for ", |
| 84 | "string at index ", i)); |
| 85 | break; |
| 86 | case CharUnit::BYTE: |
| 87 | byte_pos = AdjustedPosIndex(byte_pos, in); |
| 88 | OP_REQUIRES( |
| 89 | context, FastBoundsCheck(byte_pos, in.size() + 1), |
| 90 | errors::InvalidArgument("pos ", pos, " out of range for ", |
| 91 | "string b'", in, "' at index ", i)); |
| 92 | } |
| 93 | StringPiece sub_in = in.substr(byte_pos, byte_len); |
| 94 | output(i).assign(sub_in.data(), sub_in.size()); |
| 95 | } |
| 96 | } else { |
| 97 | // Perform Op element-wise with tensor pos/len |
| 98 | auto pos_flat = pos_tensor.flat<T>(); |
| 99 | auto len_flat = len_tensor.flat<T>(); |
| 100 | for (size_t i = 0; i < input_tensor.NumElements(); ++i) { |
| 101 | StringPiece in(input(i)); |
| 102 | const T pos = tensorflow::internal::SubtleMustCopy(pos_flat(i)); |
| 103 | const T len = tensorflow::internal::SubtleMustCopy(len_flat(i)); |
| 104 | T byte_pos = pos; |
nothing calls this directly
no test coverage detected