Home
last modified time | relevance | path

Searched refs:slice_limits (Results 1 – 24 of 24) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dslice_sinker_test.cc78 EXPECT_THAT(slice0->slice_limits(), ElementsAre(2, 9)); in TEST_F()
81 EXPECT_THAT(slice1->slice_limits(), ElementsAre(8, 9)); in TEST_F()
118 EXPECT_THAT(slice0->slice_limits(), ElementsAre(2, 9)); in TEST_F()
121 EXPECT_THAT(slice1->slice_limits(), ElementsAre(8, 9)); in TEST_F()
124 EXPECT_THAT(slice2->slice_limits(), ElementsAre(8, 4)); in TEST_F()
167 EXPECT_THAT(slice0->slice_limits(), ElementsAre(2, 9)); in TEST_F()
170 EXPECT_THAT(slice1->slice_limits(), ElementsAre(8, 9)); in TEST_F()
173 EXPECT_THAT(slice2->slice_limits(), ElementsAre(8, 2)); in TEST_F()
176 EXPECT_THAT(slice3->slice_limits(), ElementsAre(8, 9)); in TEST_F()
213 EXPECT_THAT(slice0->slice_limits(), ElementsAre(2, 9)); in TEST_F()
[all …]
Dall_to_all_decomposer.cc96 DimensionVector slice_limits(all_to_all->shape().dimensions().begin(), in ExpandInstruction() local
98 slice_limits[split_dim] = split_size; in ExpandInstruction()
107 slice_shape, operand, slice_starts, slice_limits, slice_strides))); in ExpandInstruction()
109 slice_starts[split_dim] = slice_limits[split_dim]; in ExpandInstruction()
110 slice_limits[split_dim] += split_size; in ExpandInstruction()
Dtopk_rewriter.cc117 if (has_batch && slice->slice_limits(batch_dim) != in Run()
124 k = slice->slice_limits(sort_dim); in Run()
125 } else if (k != slice->slice_limits(sort_dim)) { in Run()
Ddynamic_index_splitter_test.cc65 EXPECT_EQ(slice->slice_limits(0), i + 1); in TEST_F()
103 EXPECT_EQ(slice->slice_limits(0), i + 1); in TEST_F()
Dslice_sinker.cc40 slice_1->slice_limits() == slice_2->slice_limits() && in SameSliceConfiguration()
Dhlo_sharding_util.cc626 slice_limits(hlo.shape().rank()); in GatherEffectiveOutputSharding() local
629 slice_limits[i] = hlo.sharding().tile_assignment().dim(i); in GatherEffectiveOutputSharding()
631 slice_limits[i] = 1; in GatherEffectiveOutputSharding()
635 hlo.sharding().tile_assignment().Slice(slice_starts, slice_limits); in GatherEffectiveOutputSharding()
737 std::vector<int64> slice_starts(index_rank, 0LL), slice_limits(index_rank); in ScatterEffectiveIndexSharding() local
740 slice_limits[i] = index_sharding.tile_assignment().dim(i); in ScatterEffectiveIndexSharding()
742 slice_limits[i] = 1; in ScatterEffectiveIndexSharding()
746 index_sharding.tile_assignment().Slice(slice_starts, slice_limits); in ScatterEffectiveIndexSharding()
Dhlo_constant_folding_test.cc156 const int64 slice_limits[] = {10, 8, 6, 5, 9}; in TEST_F() local
165 shape, literal_instruction, slice_start, slice_limits, slice_strides)); in TEST_F()
Dalgebraic_simplifier.cc1252 int64 slice_end = operands[i]->slice_limits(concatenate_dimension); in HandleConcatenate()
1276 slice_end = operands[j]->slice_limits(concatenate_dimension); in HandleConcatenate()
1285 auto new_limit_indices = operands[i]->slice_limits(); in HandleConcatenate()
4053 slice->slice_limits(i) != reshape->shape().dimensions(i)) { in TryToReorderSliceAndReshape()
4102 std::vector<int64> new_limits = slice->slice_limits(); in TryToReorderSliceAndReverse()
4106 int64 limit = slice->slice_limits(rdim); in TryToReorderSliceAndReverse()
4154 std::vector<int64> new_limits = slice->slice_limits(); in HandleSlice()
4158 const int64 limit = slice->slice_limits(i); in HandleSlice()
4204 std::vector<int64> new_slice_limits = slice->slice_limits(); in HandleSlice()
4221 slice->slice_limits(dim) != in HandleSlice()
[all …]
Dhlo_instruction.cc311 std::vector<int64> slice_starts, slice_limits, slice_strides; in CreateFromProto() local
315 slice_limits.push_back(slice_dimensions.limit()); in CreateFromProto()
318 instruction = CreateSlice(shape, operands(0), slice_starts, slice_limits, in CreateFromProto()
3942 int64 HloInstruction::slice_limits(int64 dimension) const { in slice_limits() function in xla::HloInstruction
3943 return Cast<HloSliceInstruction>(this)->slice_limits(dimension); in slice_limits()
3946 const std::vector<int64>& HloInstruction::slice_limits() const { in slice_limits() function in xla::HloInstruction
3947 return Cast<HloSliceInstruction>(this)->slice_limits(); in slice_limits()
Dhlo_instruction.h1692 int64 slice_limits(int64 dimension) const;
1693 const std::vector<int64>& slice_limits() const;
Dalgebraic_simplifier_test.cc2153 EXPECT_EQ(slice->slice_limits(0), 6); in TEST_F()
2155 EXPECT_EQ(slice->slice_limits(1), 7); in TEST_F()
2157 EXPECT_EQ(slice->slice_limits(2), 32); in TEST_F()
2189 EXPECT_EQ(slice->slice_limits(0), 6); in TEST_F()
2191 EXPECT_EQ(slice->slice_limits(1), 7); in TEST_F()
2193 EXPECT_EQ(slice->slice_limits(2), 32); in TEST_F()
3477 EXPECT_EQ(computation->root_instruction()->slice_limits(0), dim0 - 2); in TEST_F()
3478 EXPECT_EQ(computation->root_instruction()->slice_limits(1), dim1 - 4); in TEST_F()
4971 EXPECT_EQ(root->slice_limits(0), 2); in TEST_F()
Dhlo_instructions.h802 int64 slice_limits(int64 dimension) const { return slice_limits_[dimension]; } in slice_limits() function
803 const std::vector<int64>& slice_limits() const { return slice_limits_; } in slice_limits() function
Ddynamic_dimension_inference.cc1104 hlo->slice_limits(dimension) != in HandleSlice()
Dhlo_verifier.cc842 slice->slice_limits(), slice->slice_strides())); in HandleSlice()
Dhlo_evaluator_typed_visitor.h2441 slice->slice_limits(), slice->slice_strides()));
/external/tensorflow/tensorflow/compiler/xla/tests/
Dslice_test.cc391 std::array<int64, 2> slice_limits; member
410 Slice(a, spec.slice_starts, spec.slice_limits, spec.slice_strides); in XLA_TEST_P()
415 input, spec.slice_starts, spec.slice_limits, spec.slice_strides); in XLA_TEST_P()
476 std::array<int64, 4> slice_limits; member
485 "__limits_", absl::StrJoin(spec.slice_limits, "x"), in R4SpecToString()
497 values, spec.slice_starts, spec.slice_limits, spec.slice_strides); in Run()
504 Slice(parameter, spec.slice_starts, spec.slice_limits, spec.slice_strides); in Run()
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/IR/
Dhlo_ops.cc123 SmallVector<int64_t, 4> slice_limits; in BuildSliceLimits() local
127 slice_limits.push_back(start_index + slice_size); in BuildSliceLimits()
129 return GetI64ElementsAttr(slice_limits, builder); in BuildSliceLimits()
1437 DenseIntElementsAttr slice_limits = BuildSliceLimits( in matchAndRewrite() local
1442 slice_limits, slice_strides); in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc1204 if (slice->slice_limits(dim) != in GetKValueInTopKWhenPartitionSortDim()
1212 k = slice->slice_limits(sort_dim); in GetKValueInTopKWhenPartitionSortDim()
1213 } else if (k != slice->slice_limits(sort_dim)) { in GetKValueInTopKWhenPartitionSortDim()
Dspmd_partitioner.cc1596 dim->set_padding_high(hlo->slice_limits(i) - in HandleSlice()
1612 std::vector<int64> limit_indices = hlo->slice_limits(); in HandleSlice()
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dhlo_function_importer.cc482 ConvertDimensions(instruction->slice_limits()), in ImportInstructionImpl()
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dmhlo_to_lhlo_with_xla.cc1235 slice.limit_indicesAttr(builder_.getI64TensorAttr(hlo_slice->slice_limits())); in EmitSliceOp()
Dlegalize_tf.cc298 auto slice_limits = llvm::to_vector<4>(type.getShape()); in SliceInMinorDims() local
300 slice_limits.begin() + major_dims); in SliceInMinorDims()
304 GetI64ElementsAttr(slice_limits, builder), in SliceInMinorDims()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dir_emitter.cc1903 ? slice->slice_limits(memcpy_dim) - slice->slice_starts(memcpy_dim) in HandleSlice()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter_unnested.cc5767 index.GetConstantWithIndexType(slice->slice_limits(dim))); in EmitElementForInputFusibleSlices()