Home
last modified time | relevance | path

Searched refs:start_indices (Results 1 – 25 of 94) sorted by relevance

1234

/external/tensorflow/tensorflow/compiler/xla/tests/
Dgather_operation_test.cc33 Literal* start_indices) { in RunTest() argument
34 RunTest(hlo_text, {operand, start_indices}); in RunTest()
63 Literal start_indices = LiteralUtil::CreateR1<int32>({0, 2}); in XLA_TEST_F() local
64 RunTest(hlo_text, &operand, &start_indices); in XLA_TEST_F()
87 Literal start_indices = LiteralUtil::CreateR1<int32>({1, 0}); in XLA_TEST_F() local
88 RunTest(hlo_text, &operand, &start_indices); in XLA_TEST_F()
108 Literal start_indices = LiteralUtil::CreateR1<int32>({0, 2}); in XLA_TEST_F() local
109 RunTest(hlo_text, &operand, &start_indices); in XLA_TEST_F()
129 Literal start_indices = LiteralUtil::CreateR2<int32>({{0, 2}, {2, 1}}); in XLA_TEST_F() local
130 RunTest(hlo_text, &operand, &start_indices); in XLA_TEST_F()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/
Dgather_expander.cc30 HloInstruction* start_indices, int64 index_vector_dim) { in TransposeIndexVectorDimToLast() argument
31 const Shape& start_indices_shape = start_indices->shape(); in TransposeIndexVectorDimToLast()
34 return start_indices; in TransposeIndexVectorDimToLast()
38 return start_indices; in TransposeIndexVectorDimToLast()
49 return MakeTransposeHlo(start_indices, permutation); in TransposeIndexVectorDimToLast()
57 HloInstruction* start_indices, int64 index_vector_dim) { in CanonicalizeGatherIndices() argument
61 TransposeIndexVectorDimToLast(start_indices, index_vector_dim)); in CanonicalizeGatherIndices()
63 index_vector_dim == start_indices->shape().dimensions_size(); in CanonicalizeGatherIndices()
156 HloInstruction* const start_indices = incoming_loop_state[1]; in GatherLoopBody() local
159 bool has_scalar_indices = start_indices->shape().dimensions_size() == 1; in GatherLoopBody()
[all …]
Dhlo_creation_utils.cc82 absl::Span<const int64> start_indices, in MakeSliceHlo() argument
87 operand->shape(), start_indices, in MakeSliceHlo()
90 slice_shape, operand, start_indices, limit_indices, strides)); in MakeSliceHlo()
136 HloInstruction* operand, absl::Span<HloInstruction* const> start_indices, in MakeDynamicSliceHlo() argument
140 start_indices.size(), in MakeDynamicSliceHlo()
141 ShapeUtil::MakeShape(start_indices[0]->shape().element_type(), {})); in MakeDynamicSliceHlo()
147 dynamic_slice_shape, operand, start_indices, slice_sizes)); in MakeDynamicSliceHlo()
151 HloInstruction* operand, HloInstruction* start_indices, in MakeDynamicSliceHlo() argument
154 CHECK_EQ(computation, start_indices->parent()); in MakeDynamicSliceHlo()
155 int64 rank = start_indices->shape().dimensions(0); in MakeDynamicSliceHlo()
[all …]
Dhlo_creation_utils.h56 absl::Span<const int64> start_indices,
88 HloInstruction* operand, absl::Span<HloInstruction* const> start_indices,
91 HloInstruction* operand, HloInstruction* start_indices,
99 HloInstruction* start_indices);
Dall_gather_decomposer.cc65 std::vector<HloInstruction*> start_indices(ag->shape().rank(), zero_index); in DecomposeAllGather() local
122 start_indices[ag->all_gather_dimension()] = in DecomposeAllGather()
128 zero->shape(), zero, ag->mutable_operand(0), start_indices)); in DecomposeAllGather()
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dsorting.cc63 std::vector<int64> start_indices(input_shape.dimensions_size(), 0); in TopK() local
68 XlaOp values = Slice(GetTupleElement(sort_result, 0), start_indices, in TopK()
70 XlaOp indices = Slice(GetTupleElement(sort_result, 1), start_indices, in TopK()
126 std::vector<int64> start_indices(input_shape.dimensions_size(), 0); in TopKWithPartitions() local
130 start_indices[last_dim] = 0; in TopKWithPartitions()
132 values = Slice(GetTupleElement(sort_result, 0), start_indices, in TopKWithPartitions()
134 indices = Slice(GetTupleElement(sort_result, 1), start_indices, in TopKWithPartitions()
141 std::vector<int64> start_indices(input_shape.dimensions_size(), 0); in TopKWithPartitions() local
144 start_indices[last_dim] = 0; in TopKWithPartitions()
147 XlaOp sliced_input = Slice(input, start_indices, limit_indices, strides); in TopKWithPartitions()
[all …]
/external/tensorflow/tensorflow/lite/kernels/internal/
Dstrided_slice_logic.h49 p->start_indices[i + pad_count] = p->start_indices[i]; in StridedSlicePadIndices()
53 p->start_indices[i] = 0; in StridedSlicePadIndices()
78 const auto* start_indices = params.start_indices; in StartForAxis() local
85 int start = start_indices[axis]; in StartForAxis()
184 const std::vector<int>& start_indices, const std::vector<int>& stop_indices, in BuildStridedSliceParams() argument
187 const int dims_count = start_indices.size(); in BuildStridedSliceParams()
193 op_params.start_indices[i] = start_indices[i]; in BuildStridedSliceParams()
Dstrided_slice_logic_test.cc35 op_params.start_indices[i] = begin.begin()[i]; in RunStridedSlicePadIndices()
43 EXPECT_EQ(op_params.start_indices[i], expected_begin.begin()[i]); in RunStridedSlicePadIndices()
/external/tensorflow/tensorflow/core/kernels/
Dstrided_slice_op.h34 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& start_indices, in operator()
42 start_i[i] = start_indices[i]; in operator()
50 input.stridedSlice(start_indices, stop_indices, strides); in operator()
82 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& start_indices,
91 start_i[i] = start_indices[i];
98 output.stridedSlice(start_indices, stop_indices, strides).device(d) =
108 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& start_indices,
116 start_i[i] = start_indices[i];
123 output.stridedSlice(start_indices, stop_indices, strides).device(d) =
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dstack_ops.cc148 std::vector<xla::XlaOp> start_indices(elem_shape.dims() + 1, in Compile() local
150 start_indices[0] = index; in Compile()
160 b, {xla::DynamicUpdateSlice(ta, update, start_indices), in Compile()
204 std::vector<xla::XlaOp> start_indices(stack_shape.dims(), in Compile() local
206 start_indices[0] = index; in Compile()
213 xla::XlaOp read = xla::DynamicSlice(ta, start_indices, slice_shape); in Compile()
Dtensor_list_utils.cc380 std::vector<xla::XlaOp> start_indices( in ExecuteTensorListPushBack() local
383 start_indices[0] = push_index; in ExecuteTensorListPushBack()
387 xla::DynamicUpdateSlice(list_part, element_part, start_indices); in ExecuteTensorListPushBack()
397 std::vector<xla::XlaOp> start_indices(element_shape.dimensions_size() + 1, in ExecuteTensorListPushBack() local
399 start_indices[0] = push_index; in ExecuteTensorListPushBack()
403 xla::DynamicUpdateSlice(list_part, update, start_indices); in ExecuteTensorListPushBack()
436 std::vector<xla::XlaOp> start_indices(list_part_shape.dimensions_size(), in ExecuteTensorListPopBack() local
438 start_indices[0] = push_index; in ExecuteTensorListPopBack()
445 xla::XlaOp read = xla::DynamicSlice(list_part, start_indices, slice_shape); in ExecuteTensorListPopBack()
484 std::vector<xla::XlaOp> start_indices(element_shape.dimensions_size() + 1, in ExecuteTensorListSetItem() local
[all …]
Dtensor_array_ops.cc125 absl::Span<const xla::XlaOp> start_indices, in DynamicAddSlice() argument
127 xla::XlaOp current = xla::DynamicSlice(operand, start_indices, update_dims); in DynamicAddSlice()
130 return xla::DynamicUpdateSlice(operand, sum, start_indices); in DynamicAddSlice()
215 std::vector<xla::XlaOp> start_indices(elem_shape.dims() + 1, in Compile() local
217 start_indices[0] = index; in Compile()
226 start_indices, dtype_); in Compile()
231 written = xla::DynamicUpdateSlice(ta, update, start_indices); in Compile()
266 std::vector<xla::XlaOp> start_indices(ta_shape.dims(), in Compile() local
268 start_indices[0] = index; in Compile()
273 xla::XlaOp read = xla::DynamicSlice(ta, start_indices, slice_shape); in Compile()
[all …]
Dunpack_op.cc67 std::vector<int64> start_indices(input_shape.dims(), 0); in Compile() local
75 start_indices[axis] = i; in Compile()
77 auto slice = xla::Slice(input, start_indices, limit_indices, strides); in Compile()
Droll_op.cc83 std::vector<xla::XlaOp> start_indices( in Compile() local
85 start_indices[cur_axis] = axis_size - offset; in Compile()
87 xla::DynamicSlice(concat, start_indices, input_shape.dim_sizes()); in Compile()
Dbatchtospace_op.cc125 std::vector<int64> start_indices(input_rank, 0); in BatchToSpace() local
133 start_indices[1 + i] = crop_start; in BatchToSpace()
136 ctx, start_indices[1 + i] <= end_indices[1 + i], in BatchToSpace()
142 xla::Slice(reshaped_permuted, start_indices, end_indices, strides); in BatchToSpace()
Ddynamic_slice_ops.cc112 xla::XlaOp start_indices = ctx->Input("start_indices"); in Compile() local
114 ctx->Input("input"), SliceVector(start_indices, rank), size_indices); in Compile()
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/
Dlegalize_gather_to_torch_index_select.cc33 auto start_indices = gather.start_indices(); in matchAndRewrite() local
34 auto start_indices_ty = start_indices.getType().cast<ShapedType>(); in matchAndRewrite()
120 operand, gather.start_indices(), rewriter.getI64IntegerAttr(0), in matchAndRewrite()
Doptimize_mhlo.cc64 !gather.start_indices().getType().cast<ShapedType>().hasRank() || in matchAndRewrite()
65 !gather.start_indices().getType().cast<ShapedType>().hasStaticShape()) { in matchAndRewrite()
121 auto gather_start_indices = gather.start_indices(); in matchAndRewrite()
/external/tensorflow/tensorflow/lite/toco/graph_transformations/
Dresolve_constant_strided_slice.cc40 int num_input_axes = op.start_indices.size(); in StridedSlice()
41 CHECK_EQ(num_input_axes, op.start_indices.size()); in StridedSlice()
57 op.begin_mask, op.end_mask, op.shrink_axis_mask, op.start_indices, in StridedSlice()
126 if (op->start_indices.empty() || op->stop_indices.empty() || in Run()
Dresolve_strided_slice_attributes.cc50 if (!op->start_indices.empty()) { in Run()
118 op->start_indices = start_array.GetBuffer<ArrayDataType::kInt32>().data; in Run()
/external/tensorflow/tensorflow/lite/toco/
Dmodel.h1108 std::vector<int> start_indices;
1123 start_indices = other.start_indices;
1136 CHECK_GE(dim_count, start_indices.size());
1137 CHECK_EQ(start_indices.size(), stop_indices.size());
1140 for (int i = start_indices.size(); i < dim_count; i++) {
1141 start_indices.push_back(0);
1150 CHECK_EQ(start_indices.size(), stop_indices.size());
1153 std::reverse(start_indices.begin(), start_indices.end());
1158 (32 - start_indices.size());
1161 (32 - start_indices.size());
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_conv_padding_legalization.cc94 std::vector<int64> start_indices(input->shape().dimensions_size(), 0); in MaybePaddedAndSlicedInput() local
103 start_indices[dim] += -conv_window->dimensions(i).padding_low(); in MaybePaddedAndSlicedInput()
113 MakeSliceHlo(input, start_indices, limit_indices, strides).ValueOrDie(); in MaybePaddedAndSlicedInput()
351 std::vector<int64> start_indices(new_backward_conv->shape().dimensions_size(), in CanonicalizeBackwardInputConvolution() local
369 start_indices[dim] += padding_low - padding_high; in CanonicalizeBackwardInputConvolution()
379 ShapeInference::InferSliceShape(new_backward_conv->shape(), start_indices, in CanonicalizeBackwardInputConvolution()
388 start_indices, limit_indices, strides)); in CanonicalizeBackwardInputConvolution()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dlegalize_hlo.cc280 auto start_indices = slice_op.start_indices(); in matchAndRewrite() local
283 for (auto pair : llvm::zip(start_indices.getValues<APInt>(), in matchAndRewrite()
292 auto start = rewriter.create<ConstOp>(slice_op.getLoc(), start_indices); in matchAndRewrite()
310 Type start_indices_element_type = op.start_indices() in matchAndRewrite()
318 start_indices_vector.reserve(op.start_indices().size()); in matchAndRewrite()
321 for (uint64_t i = 0, e = op.start_indices().size(); i < e; ++i) { in matchAndRewrite()
328 op.getLoc(), op.start_indices()[i].getType(), op.start_indices()[i], in matchAndRewrite()
1016 Value start_indices = gather_op.start_indices(); in matchAndRewrite() local
1020 ShapedType start_indices_type = start_indices.getType().cast<ShapedType>(); in matchAndRewrite()
1030 if (failed(NormalizeIndexVector(gather_op, start_indices, in matchAndRewrite()
[all …]
/external/tensorflow/tensorflow/compiler/mlir/xla/ir/
Dmlir_hlo_builder.h180 const Shape& shape, XlaOp input, XlaOp start_indices,
242 absl::Span<const int64> start_indices,
247 const Shape& shape, XlaOp operand, absl::Span<const XlaOp> start_indices,
252 absl::Span<const XlaOp> start_indices) override;
/external/tensorflow/tensorflow/compiler/xla/
Dliteral.cc890 const Shape& result_shape, absl::Span<const int64> start_indices) const { in SliceInternal()
896 new_indices[i] = indices[i] + start_indices[i]; in SliceInternal()
903 int64 dynamic_size = GetDynamicSize(dnum) - start_indices[dnum]; in SliceInternal()
912 Literal LiteralBase::Slice(absl::Span<const int64> start_indices, in Slice() argument
918 CHECK_GE(start_indices[dnum], 0); in Slice()
921 int64 dimension = limit_indices[dnum] - start_indices[dnum]; in Slice()
931 return SliceInternal<bool>(result_shape, start_indices); in Slice()
933 return SliceInternal<uint8>(result_shape, start_indices); in Slice()
935 return SliceInternal<uint16>(result_shape, start_indices); in Slice()
937 return SliceInternal<uint32>(result_shape, start_indices); in Slice()
[all …]

1234