return signalPassFailure();
    }
  }
};

std::unique_ptr<Pass> createUnfuseBatchNormPass() {
  return std::make_unique<UnfuseMhloBatchNormPass>();
}

static PassRegistration<UnfuseMhloBatchNormPass> pass;

}  // namespace odml

    }

    op->emitRemark("Not supported quantization data type.");
    return failure();
  }

 protected:
  tensorflow::quantization::QuantizationOptions quant_options_;
};

static PassRegistration<QuantizeWeightsPass> pass;

void QuantizeWeightsPass::runOnOperation() {
  MLIRContext* ctx = &getContext();
  auto module_op = getOperation();
  RewritePatternSet patterns(ctx);

class QuantizePass
    : public PassWrapper<QuantizePass, OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizePass)

  // Constructor used by the PassRegistration and only used by test.
  explicit QuantizePass() {
    quant_specs_.inference_type = tensorflow::DT_QINT8;
  }

  // Constructor used by manually creating the pass.

}
}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>> CreateOptimizeBatchMatmulPass() {
  return std::make_unique<OptimizeBatchMatmulPass>();
}

static PassRegistration<OptimizeBatchMatmulPass> pass;

}  // namespace TFL

      return signalPassFailure();
    }
  }
};

std::unique_ptr<Pass> createFoldBroadcastPass() {
  return std::make_unique<FoldBroadcastPass>();
}

static PassRegistration<FoldBroadcastPass> pass;

}  // namespace odml

struct QuantizePass : public impl::QuantizePassBase<QuantizePass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizePass)

  // Constructor used by the PassRegistration and only used by test.
  explicit QuantizePass() { quant_specs.inference_type = tensorflow::DT_QINT8; }

  // Constructor used by manually creating the pass.

  save_func_op.erase();
}

std::unique_ptr<OperationPass<ModuleOp>>
CreateMergeSaveFunctionOpsToMainPass() {
  return std::make_unique<MergeSaveFunctionOpsToMainPass>();
}

static PassRegistration<MergeSaveFunctionOpsToMainPass> pass([] {
  return CreateMergeSaveFunctionOpsToMainPass();
});

}  // namespace quant

    }
  }
};

std::unique_ptr<OperationPass<ModuleOp>>
CreateLegalizeTFXlaCallModuleToStablehloPass() {
  return std::make_unique<TFXlaCallModuleOpToStablehloPass>();
}

static PassRegistration<TFXlaCallModuleOpToStablehloPass> pass;

}  // namespace odml

    const bool enable_per_channel_quantization) {
  return std::make_unique<PreprocessOpPass>(op_set, quantization_method,
                                            enable_per_channel_quantization);
}

static PassRegistration<PreprocessOpPass> pass;

}  // namespace quant

        calibration_parameters.set_num_bins(512);
        calib_opts_.mutable_calibration_parameters()->CopyFrom(
            calibration_parameters);
        break;
      }
    }
  }
};

static PassRegistration<InsertCustomAggregationOpsPass> pass;

class AddCustomAggregationOp : public RewritePattern {
 public:
  // Does not take ownership of context, which must refer to a valid value that
  // outlives this object.