Home
last modified time | relevance | path

Searched refs:keys_shape (Results 1 – 7 of 7) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/llvm_ir/
Dsort_util.cc309 const Shape& keys_shape = values_arrays[0].GetShape(); in EmitSortInPlace() local
310 int64 rank = keys_shape.rank(); in EmitSortInPlace()
311 int64 dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); in EmitSortInPlace()
315 for (int64 dimension : LayoutUtil::MinorToMajor(keys_shape)) { in EmitSortInPlace()
317 dimensions_in_iteration_order[dim] = keys_shape.dimensions(dimension); in EmitSortInPlace()
324 Shape iteration_shape = ShapeUtil::MakeShape(keys_shape.element_type(), in EmitSortInPlace()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_layout_assignment.cc253 Shape keys_shape = instruction->operand(0)->shape(); in AddBackendConstraints() local
255 LayoutUtil::GetDefaultLayoutForRank(keys_shape.rank()); in AddBackendConstraints()
Dir_emitter_unnested.cc2862 const Shape& keys_shape = context.operand_shapes[0]; in EmitSortFromMlir() local
2867 TF_RET_CHECK(LayoutUtil::LayoutsInShapesEqual(keys_shape, in EmitSortFromMlir()
2870 LayoutUtil::LayoutsInShapesEqual(keys_shape, context.output_shapes[i])); in EmitSortFromMlir()
2890 uint64 dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); in EmitSortFromMlir()
2927 Shape standard_iteration_shape = keys_shape; in EmitSortFromMlir()
2939 Shape iteration_shape = keys_shape; in EmitSortFromMlir()
/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_dataflow_analysis_test.cc2567 Shape keys_shape = ShapeUtil::MakeShape(F32, {8}); in TEST_F() local
2569 HloInstruction::CreateParameter(0, keys_shape, "keys")); in TEST_F()
2571 auto* sort, MakeSortHlo(keys_shape, {keys}, -1, /*is_stable=*/false, in TEST_F()
2585 Shape keys_shape = ShapeUtil::MakeShape(F32, {8}); in TEST_F() local
2588 HloInstruction::CreateParameter(0, keys_shape, "keys")); in TEST_F()
2593 MakeSortHlo(ShapeUtil::MakeTupleShape({keys_shape, values_shape}), in TEST_F()
Dalgebraic_simplifier_test.cc3564 Shape keys_shape = ShapeUtil::MakeShape(F32, {1}); in TEST_F() local
3566 HloInstruction::CreateParameter(0, keys_shape, "keys")); in TEST_F()
3567 TF_ASSERT_OK(MakeSortHlo(keys_shape, {keys}, 0, /*is_stable=*/false, &builder, in TEST_F()
3580 Shape keys_shape = ShapeUtil::MakeShape(F32, {5, 0}); in TEST_F() local
3583 HloInstruction::CreateParameter(0, keys_shape, "keys")); in TEST_F()
3589 {keys_shape, values_shape, values_shape}), in TEST_F()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dir_emitter.cc548 Shape keys_shape = sort->keys()->shape(); in HandleSort() local
549 PrimitiveType keys_type = keys_shape.element_type(); in HandleSort()
562 LayoutUtil::LayoutsInShapesEqual(keys_shape, operand->shape())); in HandleSort()
564 keys_shape, ShapeUtil::GetSubshape(sort->shape(), shape_index))); in HandleSort()
585 ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(keys_shape); in HandleSort()
587 keys_shape.layout())[sort->sort_dimension()]; in HandleSort()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.cc2082 TF_ASSIGN_OR_RETURN(const Shape* keys_shape, GetShapePtr(operands[0])); in SortInternal()
2083 dimension = keys_shape->rank() - 1; in SortInternal()