<< " takes zero inputs.\n";
    return false;
  }

  // TODO(unda, 253431636): constrained kernel
  std::set<string> constrained_kernel = {"Diag",
                                         "DiagPart",
                                         "GatherNd",
                                         "GatherV2",
                                         "QuantizeAndDequantizeV2",

     {"All", "Any", "Min", "Max", "Mean", "Prod", "Sum"}},
          // clang-format on
          {"PWRED",
           {"ArgMax", "ArgMin", "DiagPart", "Softmax",
            "SparseSoftmaxCrossEntropyWithLogits", "LogSoftmax"}},
          {"REDUCEWINDOW",
           {"ArgMax", "ArgMin", "DiagPart", "Softmax",
            "SparseSoftmaxCrossEntropyWithLogits", "LogSoftmax"}},
          {"REDUCEWINDOWPW", {"BiasAddGrad", "LRN", "LRNGrad"}},

  // CHECK: return [[CLAMP]]
  func.return %0 : tensor<?xf32>
}

//===----------------------------------------------------------------------===//
// DiagPart
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: func @diag_part
// CHECK-SAME: %[[ARG:.*]]: tensor<4x3x4x3xf32>

  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

def TF_DiagPartOp : TF_Op<"DiagPart", [Pure]> {
  let summary = "Returns the diagonal part of the tensor.";

  let description = [{
This operation returns a tensor with the `diagonal` part
of the `input`. The `diagonal` part is computed as follows: