Searched refs:GetDiagonalMask (Results 1 – 4 of 4) sorted by relevance
37 XlaOp GetDiagonalMask(XlaOp x, int diagonal = 0);
57 XlaOp GetDiagonalMask(XlaOp x, int diagonal) { in GetDiagonalMask() function90 auto mask = GetDiagonalMask(x, k); in GetMatrixDiagonal()205 return Select(GetDiagonalMask(matrix, k), in SetMatrixDiagonal()
584 Square(Select(GetDiagonalMask(w_sliced), ZerosLike(w_sliced), w_sliced))); in ComputeToleranceComparison()
222 const auto mask = xla::GetDiagonalMask(output, diag_index); in SetMatrixDiag()