// Recreate the operation by using the wrapper's operands and return types.
    // TODO(fengliuai): copy the regions.
    OperationState state(op->getLoc(), op->getName().getStringRef(),
                         parent_op->getOperands(), parent_op->getResultTypes(),
                         op->getAttrs(), op->getSuccessors());
    Operation* inlined = builder.create(state);

    parent_op->replaceAllUsesWith(inlined);
    parent_op->erase();
  }

    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:

        if (llvm::isa<mlir::TF::TPUPartitionedOutputV2Op>(user)) {
          assert(user->use_empty());
          user->erase();
        }
      }
    }

    for (auto operand : cluster.getOperands()) {
      mlir::Operation* def = operand.getDefiningOp();
      if (operand.hasOneUse() &&
          llvm::isa_and_nonnull<mlir::TF::TPUPartitionedInputV2Op>(def)) {
        operand.dropAllUses();
        def->erase();

    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

                   [&](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) {

  // 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) {

            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()) {

    auto decomposition = mlir::cast<FlatSymbolRefAttr>(calledComputations[0]);

    auto composite = rewriter.create<mlir::stablehlo::CompositeOp>(
        op.getLoc(), op.getResultTypes(), op.getOperands(), name.str(), attrs,
        decomposition.getValue());
    rewriter.replaceOp(op, composite.getResults());
    return success();
  }
};

struct LegalizeStablehloCustomCallToCompositePass