constexpr char kDeviceAttr[] = "device";

#define GEN_PASS_DEF_DECOMPOSEREDUCEDATASETPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct DecomposeReduceDatasetPass
    : public impl::DecomposeReduceDatasetPassBase<DecomposeReduceDatasetPass> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<tf_device::TensorFlowDeviceDialect>();
  }

  void runOnOperation() override;