| /external/tensorflow/tensorflow/compiler/xla/service/llvm_ir/ |
| D | sort_util.cc | 307 const Shape& keys_shape = values_arrays[0].GetShape(); in EmitSortInPlace() local
308 int64 rank = keys_shape.rank(); in EmitSortInPlace()
309 int64 dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); in EmitSortInPlace()
313 for (int64 dimension : LayoutUtil::MinorToMajor(keys_shape)) { in EmitSortInPlace()
315 dimensions_in_iteration_order[dim] = keys_shape.dimensions(dimension); in EmitSortInPlace()
322 Shape iteration_shape = ShapeUtil::MakeShape(keys_shape.element_type(), in EmitSortInPlace()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | gpu_layout_assignment.cc | 220 Shape keys_shape = instruction->operand(0)->shape(); in AddBackendConstraints() local
222 LayoutUtil::GetDefaultLayoutForRank(keys_shape.rank()); in AddBackendConstraints()
|
| D | ir_emitter_unnested.cc | 1211 Shape keys_shape = sort->operand(0)->shape(); in HandleSort() local
1218 TF_RET_CHECK(LayoutUtil::LayoutsInShapesEqual(keys_shape, in HandleSort()
1221 keys_shape, ShapeUtil::GetSubshape(sort->shape(), shape_index))); in HandleSort()
1238 uint64 dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); in HandleSort()
1274 Shape standard_iteration_shape = keys_shape; in HandleSort()
1286 Shape iteration_shape = keys_shape; in HandleSort()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | tuple_points_to_analysis_test.cc | 1071 Shape keys_shape = ShapeUtil::MakeShape(F32, {8}); in TEST_F() local
1073 HloInstruction::CreateParameter(0, keys_shape, "keys")); in TEST_F()
1075 auto* sort, MakeSortHlo(keys_shape, {keys}, 0, /*is_stable=*/false, in TEST_F()
1089 Shape keys_shape = ShapeUtil::MakeShape(F32, {8}); in TEST_F() local
1092 HloInstruction::CreateParameter(0, keys_shape, "keys")); in TEST_F()
1097 MakeSortHlo(ShapeUtil::MakeTupleShape({keys_shape, values_shape}), in TEST_F()
|
| D | hlo_dataflow_analysis_test.cc | 2450 Shape keys_shape = ShapeUtil::MakeShape(F32, {8}); in TEST_F() local
2452 HloInstruction::CreateParameter(0, keys_shape, "keys")); in TEST_F()
2454 auto* sort, MakeSortHlo(keys_shape, {keys}, -1, /*is_stable=*/false, in TEST_F()
2468 Shape keys_shape = ShapeUtil::MakeShape(F32, {8}); in TEST_F() local
2471 HloInstruction::CreateParameter(0, keys_shape, "keys")); in TEST_F()
2476 MakeSortHlo(ShapeUtil::MakeTupleShape({keys_shape, values_shape}), in TEST_F()
|
| D | algebraic_simplifier_test.cc | 2804 Shape keys_shape = ShapeUtil::MakeShape(F32, {1}); in TEST_F() local
2806 HloInstruction::CreateParameter(0, keys_shape, "keys")); in TEST_F()
2807 TF_ASSERT_OK(MakeSortHlo(keys_shape, {keys}, 0, /*is_stable=*/false, &builder, in TEST_F()
2820 Shape keys_shape = ShapeUtil::MakeShape(F32, {5, 0}); in TEST_F() local
2823 HloInstruction::CreateParameter(0, keys_shape, "keys")); in TEST_F()
2829 {keys_shape, values_shape, values_shape}), in TEST_F()
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | ir_emitter.cc | 507 Shape keys_shape = sort->keys()->shape(); in HandleSort() local
508 PrimitiveType keys_type = keys_shape.element_type(); in HandleSort()
537 LayoutUtil::LayoutsInShapesEqual(keys_shape, operand->shape())); in HandleSort()
539 keys_shape, ShapeUtil::GetSubshape(sort->shape(), shape_index))); in HandleSort()
560 ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(keys_shape); in HandleSort()
562 keys_shape.layout())[sort->sort_dimension()]; in HandleSort()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.cc | 1698 TF_ASSIGN_OR_RETURN(const Shape& keys_shape, GetShape(operands[0])); in Sort()
1699 dimension = keys_shape.rank() - 1; in Sort()
|