Function InternalCompute

tensorflow/core/kernels/regex_replace_op.cc:32–58

Execute the specified regex using the given context. Context requirements: - "input" string Tensor at input_index=0 - "output" string Tensor at output_index=0

//  - "input" string Tensor at input_index=0
//  - "output" string Tensor at output_index=0
Status InternalCompute(const RE2& match, const string& rewrite,
                       const bool replace_global, OpKernelContext* ctx) {
  const Tensor* input_tensor;
  TF_RETURN_IF_ERROR(ctx->input("input", &input_tensor));
  Tensor* output_tensor;
  std::unique_ptr<Tensor> maybe_forwarded =
      ctx->forward_input(0 /*input_index*/, 0 /*output_index*/,
                         tensorflow::DT_STRING, input_tensor->shape(),
                         ctx->input_memory_type(0), ctx->input_alloc_attr(0));
  if (maybe_forwarded) {
    output_tensor = maybe_forwarded.get();
    TF_RETURN_IF_ERROR(ctx->set_output("output", *output_tensor));
  } else {
    TF_RETURN_IF_ERROR(
        ctx->allocate_output("output", input_tensor->shape(), &output_tensor));
    output_tensor->flat<tstring>() = input_tensor->flat<tstring>();
  }
  auto output_flat = output_tensor->flat<tstring>();
  for (size_t i = 0; i < output_flat.size(); ++i) {
    if (replace_global) {
      RE2::GlobalReplace(&output_flat(i), match, rewrite);
    } else {
      RE2::Replace(&output_flat(i), match, rewrite);
    }
  }
  return Status::OK();
}
}  // namespace

class RegexReplaceOp : public OpKernel {