Searched refs:diag_blocks (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | triangular_solve_expander.cc | 47 XlaOp diag_blocks; in DiagonalBlocks() local
86 diag_blocks = Gather(a, start_indices, dim_numbers, slice_sizes); in DiagonalBlocks()
113 diag_blocks = in DiagonalBlocks()
114 ConcatInDim(builder, {diag_blocks, last_blocks}, ndims - 2); in DiagonalBlocks()
116 diag_blocks = last_blocks; in DiagonalBlocks()
120 return diag_blocks; in DiagonalBlocks()
124 XlaOp InvertDiagonalBlocks(XlaOp diag_blocks, bool lower, bool transpose_a, in InvertDiagonalBlocks() argument
127 XlaBuilder* builder = diag_blocks.builder(); in InvertDiagonalBlocks()
131 TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(diag_blocks)); in InvertDiagonalBlocks()
135 diag_blocks = Reshape(diag_blocks, {num_blocks, block_size, block_size}); in InvertDiagonalBlocks()
[all …]
|