int32_t k_indices = k_indices_or.value();
    // We don't match two SliceOp with different sizes.
    if (k_values != k_indices && !values.use_empty() && !indices.use_empty())
      return failure();

    // Start replacing.
    auto k = !values.use_empty() ? k_values : k_indices;
    // Build scalar tensor k.
    auto k_ty = mlir::RankedTensorType::get({}, rewriter.getIntegerType(32));

    if (failed(UpdateReadUses(read_var, partitioned_input, partitioned_read,
                              partitioned_reads, is_packed)))
      return failure();
    read_var->erase();
    if (partitioned_input->use_empty()) partitioned_input->erase();
  }
  return mlir::TF::RemoveSingletonParallelExecuteOp(parallel_execute, &builder);
}

void TPUResourceReadsWritesPartitioningPass::runOnOperation() {

  for (int i = 0; i < expected_num_outputs; ++i) {
    auto it = expected_consumer_indices.find(i);
    if (it == expected_consumer_indices.end()) {
      // Unexpected consumer.
      if (!call_op->getResult(i).use_empty()) return failure();
    }
  }
  return success();
}

LogicalResult CheckFusableKerasLstm(func::FuncOp lstm_func, ModuleOp module) {
  for (auto func : module.getOps<func::FuncOp>()) {

    if (op.getOperation()->template hasTrait<OpTrait::IsTerminator>()) {
      return failure();
    }
    for (auto result : op.getOperation()->getOpResults()) {
      if (!result.use_empty()) {
        return failure();
      }
    }
    // Remove if the custom op is in the provided map and is NoSideEffect.
    auto custom_op = llvm::isa<CustomOp>(op);
    if (custom_op) {

                                arith::ConstantOp op) const {
    for (auto connected_op : op.getResult().getUsers()) {
      auto cast_op = dyn_cast_or_null<CastOp>(connected_op);
      if (!cast_op || cast_op.getResult().use_empty()) continue;

      // Get types
      Type old_result_type = op.getResult().getType();
      ShapedType new_result_type =
          mlir::dyn_cast<ShapedType>(cast_op.getType());

  for (auto [src_arg_idx, src_arg] :
       llvm::enumerate(src_func_op.getArguments())) {
    // No need to create a mapping when there is no usage - it will not affect
    // the cloning.
    if (src_arg.use_empty()) continue;

    const unsigned main_arg_idx = main_func_op.getNumArguments();

    const DictionaryAttr main_arg_attr =
        src_func_op.getArgAttrDict(src_arg_idx);

  OpBuilder builder(main_func.getContext());
  for (auto constant_op :
       main_func.getBody().getOps<mlir::stablehlo::ConstantOp>()) {
    builder.setInsertionPointAfter(constant_op);
    if (constant_op.getResult().use_empty() ||
        constant_op.getResult().hasOneUse())
      continue;
    // Do not duplicate constant op if the size is too large.
    // 32 is chosen to be larger than all constants useful for shape references,

    return signalPassFailure();

  // Eliminate TPUCompilationResultOp now that the rewrite is complete.
  for (auto& it : compilation_results) {
    for (auto op : it.second) {
      if (!op.use_empty()) {
        mlir::InFlightDiagnostic err = op.emitError("uses remain post rewrite");
        for (auto user : op->getUsers())
          err.attachNote(user->getLoc()) << "remaining user";
        return signalPassFailure();

  }

  // Remove the read ops if they have no more uses.
  for (const auto& entry : var_access_info.per_resource_info) {
    const auto& info = entry.getSecond();
    if (info.read->use_empty()) info.read->erase();
  }
  return success();
}

// Checks if an ops parent is a tf_device.parallel_execute and the region the
// op is in is perfectly wrapped.

      then_name, else_name, if_region.getIsStateless());
  CopyAndOverrideAttributes(if_region, if_op, &builder);

  if_region.replaceAllUsesWith(if_op.getResults());
  if_region.erase();

  if (to_bool && to_bool.use_empty()) to_bool.erase();
  return success();
}

// Transform CaseRegion to Case.
LogicalResult RegionControlFlowToFunctional::ConvertCaseOp(
    SymbolTableCollection& symbol_table, CaseRegionOp case_region) {