Searched refs:sort_dimension (Results 1 – 9 of 9) sorted by relevance
54 sort->sort_dimension() && in ExpandInstruction()69 if (iota_shape.dimensions(sort->sort_dimension()) > in ExpandInstruction()76 HloInstruction::CreateIota(iota_shape, sort->sort_dimension())); in ExpandInstruction()
85 iota->iota_dimension() != sort->sort_dimension()) { in Run()91 const int64 sort_dim = sort->sort_dimension(); in Run()
637 int64 sort_dimension() const { return dimensions(0); } in sort_dimension() function
1198 int64 sort_dim = sort->sort_dimension(); in RewriteDynamicSort()
1169 iota->iota_dimension() != sort->sort_dimension()) { in GetKValueInTopKWhenPartitionSortDim()1173 const int64 sort_dim = sort->sort_dimension(); in GetKValueInTopKWhenPartitionSortDim()
1658 const int64 sort_dim = sort->sort_dimension(); in HandleSort()
496 builder_->getI64IntegerAttr(sort_instruction->sort_dimension()), in ImportInstructionImpl()
420 sort.dimensionAttr(builder_.getI64IntegerAttr(sort_instr->sort_dimension())); in EmitSortOp()
587 keys_shape.layout())[sort->sort_dimension()]; in HandleSort()