}
  return success();
}

void RewriteTPUEmbeddingOps::runOnOperation() {
  func::FuncOp func = getOperation();
  if (failed(RunOnRegion(&func.getBody()))) return signalPassFailure();

  func.walk([&](Operation* op) {
    for (Region& region : op->getRegions()) {
      if (failed(RunOnRegion(&region))) return signalPassFailure();
    }
  });
}

           << replicate.getN() << ")";
  }

  // add block arg to region. This is in $body.
  unsigned block_arg_idx = replicate.GetNumReplicatedBlockArguments();
  Block& block = replicate.GetBody();
  block_arg = block.insertArgument(block_arg_idx, type, replicate.getLoc());

  // add to $replicated_inputs. This also updates OperandSegmentSizeAttr.
  replicate.getReplicatedInputsMutable().append(inputs);

  // Create new function and extract region body into the function.
  auto outlined_func = builder.create<func::FuncOp>(loc, name, type);
  Region& func_region = outlined_func.getBody();
  func_region.takeBody(region);
  Block& first_block = func_region.front();

  // Replace all external uses with function arguments.
  for (auto it : llvm::enumerate(extern_values)) {

  // Create the ControlNodeOp.
  auto ctrl_op = op_builder.create<mlir::TFL::ControlNodeOp>(
      op_loc, cloned_op->getResultTypes(),
      mlir::TFL::ControlType::get(op->getContext()), control_tokens);
  ctrl_op.getBody().takeBody(region);

  // Store the control_token output for use by downstream nodes.
  maybe_control_node->second.outgoing = ctrl_op.getControl();

  // Remove the original op.

  if (result_tuple_ty.size() != 1) return;
  auto result_ty = mlir::dyn_cast<ShapedType>(result_tuple_ty.getType(0));

  func.eraseBody();
  func.addEntryBlock();
  ImplicitLocOpBuilder builder(func.getLoc(), func.getBody());
  llvm::SmallVector<int32_t> result_shape_i32;
  auto result_shape = result_ty.getShape();
  for (auto element : result_shape) {
    result_shape_i32.push_back(static_cast<int32_t>(element));
  }

}

// Get the reads to hoist in the `replicate`.
LogicalResult GetReads(FuncOp func, ReplicateOp replicate,
                       llvm::SmallVector<ReadVariableOp, 4>& reads) {
  for (Operation& op : replicate.getBody().front()) {
    if (auto read = llvm::dyn_cast<ReadVariableOp>(&op)) {
      bool is_cpu_read;
      if (failed(IsCpuRead(func, read, is_cpu_read))) return failure();
      if (is_cpu_read) reads.push_back(read);
    }

    bool enable_per_channel_quantization) {
  // A map to find an attribute from its identifier.
  llvm::StringMap<Attribute> identifier_to_attr;

  for (Operation& inner_op : float_func.getBody().front().getOperations()) {
    if (!inner_op.hasAttr(kAttrMapAttribute)) continue;
    // Insert quantization related attribute if they exists. Quantization

			t.Errorf("test[%d]: Body = nil", i)
			continue
		}
		if req.GetBody == nil {
			t.Errorf("test[%d]: GetBody = nil", i)
			continue
		}
		slurp1, err := io.ReadAll(req.Body)
		if err != nil {
			t.Errorf("test[%d]: ReadAll(Body) = %v", i, err)
		}
		newBody, err := req.GetBody()
		if err != nil {
			t.Errorf("test[%d]: GetBody = %v", i, err)
		}
		slurp2, err := io.ReadAll(newBody)

  }
  return success();
}

void SplitMergedOperandsPass::runOnOperation() {
  llvm::DenseSet<Value> stateful_values;
  auto func = getOperation();
  OpBuilder builder(func);
  for (auto& bb : func.getBody()) {
    for (auto& op : bb) {
      if (failed(DuplicateValueIfNeeded(&op, &stateful_values, &builder))) {
        func.emitError() << "Failed to duplicate values for the stateful op\n";
        return signalPassFailure();

                                new_while_op.getCond().end());
    rewriter.inlineRegionBefore(while_op.getBody(), new_while_op.getBody(),
                                new_while_op.getBody().end());

    auto &new_body_block = new_while_op.getBody().front();
    auto &new_cond_block = new_while_op.getCond().front();
    auto &new_body_yield = *new_body_block.getTerminator();