assert(mlir::isa<mlir::tf_device::ReturnOp>(return_op));

  Value old_val, new_val;
  for (auto it : llvm::zip(cluster_op.getResults(), return_op.getOperands())) {
    std::tie(old_val, new_val) = it;
    old_val.replaceAllUsesExcept(new_val, &return_op);
  }

  cluster_op.erase();
}

void XlaInlineDeviceOpsPass::runOnOperation() {
  ModuleOp module = getOperation();

  llvm::SmallVector<tf_device::ClusterOp, 4> ops;

        op_result.setType(ToBfloat16Type(original_type));
        const Value cast =
            builder.create<TF::CastOp>(op->getLoc(), original_type, op_result);
        op_result.replaceAllUsesExcept(cast, cast.getDefiningOp());
      }
    }
    return WalkResult::advance();
  });

  if (result.wasInterrupted()) return signalPassFailure();
}

}  // namespace

    // Wire input/output nodes.
    new_dequantize_op->setOperand(0, user_op->getResult(0));
    new_dequantize_op->getResult(0).setType(user_op->getResult(0).getType());
    user_op->getResult(0).replaceAllUsesExcept(new_dequantize_op->getResult(0),
                                               new_dequantize_op);
    user_op->getResult(0).setType(new_user_op_type);
  }

  LogicalResult matchAndRewrite(TF::PartitionedCallOp op,

  auto type = result.getType();
  result.setType(converter.convertType(result.getType()));
  auto new_value = converter.materializeArgumentConversion(
      rewriter, result.getLoc(), type, {result});
  rewriter.replaceAllUsesExcept(result, new_value, new_value.getDefiningOp());
}

// Wrap operands in an an unrealized cast to create a cast to buffer any type
// changes to the operand, and apply type converter to operands:

      Operation *aggregator_op = rewriter.create<TF::CustomAggregatorOp>(
          op->getLoc(), output_types, value, attributes);

      Value aggregator_op_result = aggregator_op->getOpResult(0);
      value.replaceAllUsesExcept(aggregator_op_result, aggregator_op);
    }

    return success();
  }

 private:
  CalibrationOptions calib_opts_;
};

void InsertCustomAggregationOpsPass::runOnOperation() {

        // 'sharding' attribute.
        auto updated_arg = builder.create<TF::XlaShardingOp>(
            func.getLoc(), arg.getType(), arg, sharding, sharding);
        func.getArgument(i).replaceAllUsesExcept(
            updated_arg, llvm::SmallPtrSet<Operation*, 1>({updated_arg}));
      }

      func.removeArgAttr(i, builder.getStringAttr(kShardingAttr));
    }
  }
}

  for (auto [original_output_value, xla_call_module_op_result_value] :
       llvm::zip_equal(outputs, xla_call_module_op->getResults())) {
    original_output_value.replaceAllUsesExcept(xla_call_module_op_result_value,
                                               /*exceptedUser=*/nullptr);
  }
}

// Replaces the StableHLO ops with a separate XlaCallModuleOp, then wires it

  Operation* terminator_op =
      old_op->getParentOfType<func::FuncOp>().front().getTerminator();
  llvm::SmallPtrSet<mlir::Operation*, 1> exceptions = {terminator_op};
  old_val.replaceAllUsesExcept(new_val, exceptions);
}

struct TensorArrayStats {
  // Whether a write op should accumulate with the old value. Set to true if
  // this is a gradient.
  bool accumulate_on_write;

    auto dq = rewriter.create<TFL::DequantizeOp>(
        op->getLoc(), /*output=*/dequantized_rhs_type,
        /*input=*/rhs);
    rewriter.replaceAllUsesExcept(rhs, dq.getOutput(), dq);
  }
};

// Splits hybrid quantized `stablehlo.convolution` into `tfl.dequantize` and
// float `stablehlo.convolution` op. Weight tensor is transposed to match the

    }

    for (auto replace_pair : replace_map) {
      Value result = replace_pair.first;
      quantfork::StorageCastOp scast_op = replace_pair.second;
      result.replaceAllUsesExcept(scast_op, scast_op);
    }

    // Make a copy of the quantized function.
    auto module = call_op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);