OpBuilder &builder) {
  llvm::SmallVector<Value> non_resource_args, resource_args;
  bool has_resources = false, in_order = true;
  for (const Value &arg : cluster_func_op.getOperands()) {
    if (!mlir::isa<TF::ResourceType>(getElementTypeOrSelf(arg.getType()))) {
      non_resource_args.push_back(arg);
      if (has_resources) in_order = false;
    } else {
      resource_args.push_back(arg);

  module.walk([&](tf_device::LaunchFuncOp launch) {
    OpBuilder builder(launch);
    auto call_op = builder.create<TF::PartitionedCallOp>(
        module.getLoc(), launch.getResultTypes(), launch.getOperands(),
        SymbolRefAttr::get(builder.getContext(), launch.getFunc()),
        /*config=*/builder.getStringAttr(""),
        /*config_proto=*/builder.getStringAttr(""),
        /*executor_type=*/builder.getStringAttr(""));

    } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) {
      for (auto &region : while_op->getRegions()) {
        for (auto [arg, value] :
             llvm::zip(region.getArguments(), while_op->getOperands())) {
          this->join(getLatticeElement(arg), *getLatticeElement(value));
        }
      }
    } else if (auto while_op = dyn_cast<TF::WhileOp>(op)) {

    for (auto operand_and_idx : llvm::enumerate(op->getOperands()))
      TF_RETURN_IF_ERROR(AddEdgeBetweenNodes(operand_and_idx.value(), dst_node,
                                             operand_and_idx.index()));

    operand_offset = op->getNumOperands();
  }

  // For all other ops (including tf_executor.island), add remaining edges.
  for (auto operand_and_idx : llvm::enumerate(inst->getOperands()))
    TF_RETURN_IF_ERROR(

    for (auto operand_and_idx : llvm::enumerate(op->getOperands()))
      TF_RETURN_IF_ERROR(AddEdgeBetweenNodes(operand_and_idx.value(), dst_node,
                                             operand_and_idx.index()));

    operand_offset = op->getNumOperands();
  }

  // For all other ops (including tf_executor.island), add remaining edges.
  for (auto operand_and_idx : llvm::enumerate(inst->getOperands()))
    TF_RETURN_IF_ERROR(

  SmallVector<NamedAttribute> attrs{op->getAttrs()};
  attrs.push_back(call_target);
  auto custom_call = rewriter.create<mlir::stablehlo::CustomCallOp>(
      op->getLoc(), op->getResultTypes(), op->getOperands(), attrs);
  rewriter.replaceOp(op, custom_call.getResults());
  return success();
}

}  // namespace

template <typename OpTy>
class SmuggleOpPattern : public OpRewritePattern<OpTy> {
 public:

    auto new_gather_op = rewriter.create<TFL::GatherOp>(
        op.getLoc(),
        /*result=*/
        mlir::cast<TensorType>(op.getResult().getType())
            .clone(builder.getI4Type()),
        /*operand=*/op.getOperands(), op->getAttrs());
    rewriter.replaceAllUsesWith(op.getResult(), new_gather_op.getResult());

    return success();
  }
};

class ReduceTypePrecisionPass

  // PartitionedCallOp does not have kQuantTraitAttrName, and therefore won't
  // get quantized.
  auto new_call_op = rewriter.create<TF::PartitionedCallOp>(
      call_op.getLoc(), call_op.getResultTypes(), call_op.getOperands(),
      FlatSymbolRefAttr::get(new_ref_func_name));
  return new_call_op;
}

Operation *DuplicateOp(TF::XlaCallModuleOp call_op, PatternRewriter &rewriter,
                       const StringAttr &new_ref_func_name) {

                   [&](Type type) { return type != resultType; })) {
    p.printGenericOp(op, /*printOpName=*/false);
    return;
  }

  p << ' ';
  p.printOperands(op->getOperands());
  p.printOptionalAttrDict(op->getAttrs());
  // Now we can output only one type for all operands and the result.
  p << " : " << resultType;
}

Operation* getDefiningBroadcastArgsOp(Value operand) {

            mlir::cast<IntegerAttr>(mirrored_index).getInt());
      }
    }
    auto func =
        llvm::cast<func::FuncOp>(m.lookupSymbol(cluster_func.getFunc()));
    for (auto entry : llvm::enumerate(cluster_func.getOperands())) {
      auto operand = SkipIdentityAndReadVariable(entry.value());
      auto block_arg = mlir::dyn_cast<BlockArgument>(operand);
      if (block_arg && block_arg.getOwner() == &replicate.GetBody()) {