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

// LayoutAssignmentPass assigns optimal data layout (data format) for all
// layout sensitive operations.
class LayoutAssignmentPass
    : public impl::LayoutAssignmentPassBase<LayoutAssignmentPass> {
 public:
  LayoutAssignmentPass() = default;
  explicit LayoutAssignmentPass(const std::string& force_data_format) {
    force_data_format_ = force_data_format;
  }

  let summary = "Device assignment in TF dialect using the device specified in the function attribute.";
  let constructor = "TF::CreateTFDeviceAssignmentByFuncAttrPass()";
}

def LayoutAssignmentPass : Pass<"tf-layout-assignment", "mlir::func::FuncOp"> {
  let summary = "Layout assignment pass.";
  let constructor = "TF::CreateLayoutAssignmentPass()";
  let options = [