Searched refs:GetMatrixDiagonalViaGather (Results 1 – 4 of 4) sorted by relevance
48 XlaOp GetMatrixDiagonalViaGather(XlaOp x, int k = 0);
109 XlaOp GetMatrixDiagonalViaGather(XlaOp x, int k) { in GetMatrixDiagonalViaGather() function
152 xla::Reshape(is_gpu_ ? xla::GetMatrixDiagonalViaGather(reshape_input) in Compile()
404 0, is_gpu_ ? xla::GetMatrixDiagonalViaGather(input, upper_diag_index) in Compile()411 is_gpu_ ? xla::GetMatrixDiagonalViaGather(input, diag_index) in Compile()