class PrepareTpuComputationForTfExportPass
    : public impl::PrepareTpuComputationForTfExportPassBase<
          PrepareTpuComputationForTfExportPass> {
  void runOnOperation() override;
};

class RewriteXlaHostComputeMlir
    : public OpRewritePattern<TF::_XlaHostComputeMlirOp> {
 public:
  using OpRewritePattern<TF::_XlaHostComputeMlirOp>::OpRewritePattern;

  LogicalResult match(TF::_XlaHostComputeMlirOp op) const override {