new_result_types.push_back(old_result.getType());

  OpBuilder builder(parallel_execute);
  auto new_parallel_execute = builder.create<tf_device::ParallelExecuteOp>(
      parallel_execute.getLoc(), num_regions, new_result_types);

  for (auto region : llvm::zip(new_parallel_execute.getRegions(),
                               parallel_execute.getRegions()))
    std::get<0>(region).takeBody(std::get<1>(region));

    mode_string_value.emplace_back(is_training ? "train" : "inference");
    builder->setInsertionPoint(embedding_op);
    auto enqueue_mode = builder->create<TF::ConstOp>(
        embedding_op->getLoc(),
        DenseStringElementsAttr::get(
            RankedTensorType::get({}, builder->getType<TF::StringType>()),
            mode_string_value));

    auto outside_compilation_attr =

      }
    }
  };

  for (int i = 0, e = func.getNumArguments(); i != e; ++i) {
    BlockArgument arg = func.getArgument(i);
    auto* arg_block = arg.getOwner();
    add_quantize_op(arg.getLoc(), arg.getType(), arg_block,
                    std::next(arg_block->begin(), i), arg, i);
  }

  return false;
}

#include "tensorflow/compiler/mlir/lite/utils/generated_op_quant_spec_getters.inc"

    if (it.second) {
      auto resource_type = op.getResult(0).getType();
      op.getResult(0).replaceAllUsesWith(
          block.addArgument(resource_type, op.getLoc()));
      AddEntryFunctionInput(
          absl::StrCat("hash_table_", it.first->getValue(), ":0"), target_func);
      // Avoid deleting the op here, clone it to the caller function first.

        }
      }

      // Replace the arg with a tf.Const op in the function body.
      builder.setInsertionPointToStart(&func.getBody().front());
      auto const_op = builder.create<TF::ConstOp>(global_tensor.getLoc(),
                                                  *global_tensor.getValue());
      args_to_erase.set(val.getArgNumber());
      for (auto read_op : read_variable_ops_to_erase) {

    auto preceding_qcast =
        q_op.getArg().getDefiningOp<quantfork::QuantizeCastOp>();
    if (!preceding_qcast) return failure();

    auto new_qcast = rewriter.create<quantfork::QuantizeCastOp>(
        q_op.getLoc(), q_op.getType(), preceding_qcast.getArg());
    new_qcast->setAttr(kVolatileOpAttrName, rewriter.getUnitAttr());
    q_op->replaceAllUsesWith(new_qcast);
    return success();
  }
};

    if (res_type != GetResourceSubtype(rng_op.getResource())) {
      return rewriter.notifyMatchFailure(op, "unexpected resource subtype");
    }

    Location loc = op->getLoc();

    // Read the state value from the resource.
    Value state =
        rewriter.create<ReadVariableOp>(loc, res_type, rng_op.getResource());

    // Extract the key and counter from the state.

  if (!status.ok()) {
    return emitRemark(op_->getLoc())
           << "failed to create XlaContext resource: " << status.ToString();
  }
  params_.step_container = step_container_.get();

  absl::StatusOr<int64_t> version_or = tensorflow::GetTfGraphProducerVersion(
      op_->getParentOfType<mlir::ModuleOp>());
  if (!version_or.ok()) {
    return emitError(op_->getLoc()) << version_or.status().ToString();
  }

    if (block_arg.getOwner() != replicate_block) return;

    OpBuilder builder(shape_op);
    auto new_shape_op = builder.create<TF::VariableShapeOp>(
        shape_op.getLoc(), shape_op.getType(),
        replicate_op.GetReplicaOperandForBlockArgument(block_arg,
                                                       /*replica=*/0));
    shape_op.replaceAllUsesWith(new_shape_op.getOperation());

  TypeRange result_types{ValueRange{results}};
  auto func_type = FunctionType::get(context, arg_types, result_types);

  SmallVector<Location> arg_locs;
  for (Value arg : arguments) {
    arg_locs.push_back(arg.getLoc());
  }

  auto wrap_func = builder.create<func::FuncOp>(location, func_name, func_type);
  wrap_func.setVisibility(SymbolTable::Visibility::Private);