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;

          } : () -> tensor<i32>
          tf_executor.fetch %outputs : tensor<i32>
        }
        return %0 : tensor<i32>
      }
    ```
  }];
}

def DecomposeReduceDatasetPass : Pass<"tf-decompose-reduce-dataset", "mlir::func::FuncOp"> {
  let summary = "Decomposes ReduceDataset op into dataset operations.";

  let description = [{