}

//===----------------------------------------------------------------------===//
// MatrixDiagPart
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: func @matrix_diag_part
// CHECK-SAME: %[[ARG:.*]]: tensor<7x140x128xi32>
func.func @matrix_diag_part(%arg0: tensor<7x140x128xi32>) -> tensor<7x22x128xi32> {

# A main diagonal from each batch.
tf.matrix_diag_part(input) ==> [[1, 6, 7],  # Output shape: (2, 3)
                                [5, 2, 7]]

# A superdiagonal from each batch.
tf.matrix_diag_part(input, k = 1)
  ==> [[2, 7, 6],  # Output shape: (2, 3)
       [4, 3, 8]]

# A band from each batch.
tf.matrix_diag_part(input, k = (-1, 2))
  ==> [[[0, 3, 8],  # Output shape: (2, 4, 3)

    "Identity",
    "IdentityN",
    "InplaceAdd",
    "InplaceSub",
    "InplaceUpdate",
    "InvertPermutation",
    "ListDiff",
    "MatrixBandPart",
    "MatrixDiag",
    "MatrixDiagPart",
    "MatrixDiagPartV2",
    "MatrixDiagPartV3",
    "MatrixDiagV2",
    "MatrixDiagV3",
    "MatrixSetDiag",
    "MatrixSetDiagV2",
    "MatrixSetDiagV3",
    "MirrorPad",
    "OneHot",

    // Replace all out-of-bounds values in the result with padding_value.
    Value result =
        rewriter.create<SelectOp>(loc, b_in_bounds, gather, b_padding);

    if (num_diags == 1) {
      // matrix_diag_part folds away the 1-sized band dimension if we only
      // extract a single diagonal.
      result = rewriter.create<ReshapeOp>(loc, op.getType(), result);
    }

    rewriter.replaceOp(op, result);

Status MatrixDiagGrad(const Scope& scope, const Operation& op,
                      const std::vector<Output>& grad_inputs,
                      std::vector<Output>* grad_outputs) {
  grad_outputs->push_back(MatrixDiagPart(scope, grad_inputs[0]));
  return scope.status();
}
REGISTER_GRADIENT_OP("MatrixDiag", MatrixDiagGrad);

Status MatrixBandPartGrad(const Scope& scope, const Operation& op,

*   Fixed bug causing incorrect number of threads to be used for multi-threaded
    benchmarks.
*   Performance optimizations for `batch_matmul` on multi-core CPUs.
*   Improve trace, `matrix_set_diag`, `matrix_diag_part` and their gradients to
    work for rectangular matrices.
*   Support for SVD of complex valued matrices.

## Thanks to our Contributors

      "L2Loss",
      "LeakyRelu",
      "LinSpace",
      "ListDiff",
      "LogMatrixDeterminant",
      "LowerBound",
      "MatMul",
      "MatrixBandPart",
      "MatrixDiag",
      "MatrixDiagPart",
      "MatrixDiagPartV2",
      "MatrixDiagPartV3",
      "MatrixDiagV2",
      "MatrixDiagV3",
      "MatrixInverse",
      "MatrixSetDiag",
      "MatrixSetDiagV2",
      "MatrixSetDiagV3",