Searched refs:diag_blocks (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | triangular_solve_expander.cc | 50 XlaOp diag_blocks; in DiagonalBlocks() local
89 diag_blocks = Gather(a, start_indices, dim_numbers, slice_sizes); in DiagonalBlocks()
123 diag_blocks = in DiagonalBlocks()
124 ConcatInDim(builder, {diag_blocks, last_blocks}, ndims - 2); in DiagonalBlocks()
126 diag_blocks = last_blocks; in DiagonalBlocks()
130 return diag_blocks; in DiagonalBlocks()
245 XlaOp diag_blocks, bool lower_triangular, in InvertDiagonalBlocks() argument
247 XlaBuilder* builder = diag_blocks.builder(); in InvertDiagonalBlocks()
251 TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(diag_blocks)); in InvertDiagonalBlocks()
255 diag_blocks = Reshape(diag_blocks, {num_blocks, block_size, block_size}); in InvertDiagonalBlocks()
[all …]
|
| D | triangular_solve_expander.h | 47 virtual XlaOp InvertDiagonalBlocks(XlaOp diag_blocks, bool lower_triangular,
|