Home
last modified time | relevance | path

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

12

/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()
84 Literal start_indices = LiteralUtil::CreateR1<int32>({0, 2}); in XLA_TEST_F() local
85 RunTest(hlo_text, &operand, &start_indices); in XLA_TEST_F()
105 Literal start_indices = LiteralUtil::CreateR2<int32>({{0, 2}, {2, 1}}); in XLA_TEST_F() local
106 RunTest(hlo_text, &operand, &start_indices); in XLA_TEST_F()
126 Literal start_indices = in XLA_TEST_F() local
128 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.cc71 absl::Span<const int64> start_indices, in MakeSliceHlo() argument
76 operand->shape(), start_indices, in MakeSliceHlo()
79 slice_shape, operand, start_indices, limit_indices, strides)); in MakeSliceHlo()
122 HloInstruction* operand, HloInstruction* start_indices, in MakeDynamicSliceHlo() argument
125 CHECK_EQ(computation, start_indices->parent()); in MakeDynamicSliceHlo()
126 int64 rank = start_indices->shape().dimensions(0); in MakeDynamicSliceHlo()
131 ShapeUtil::MakeShape(start_indices->shape().element_type(), {1}), in MakeDynamicSliceHlo()
132 start_indices, {i}, {i + 1}, {1})); in MakeDynamicSliceHlo()
135 ShapeUtil::MakeShape(start_indices->shape().element_type(), {}), in MakeDynamicSliceHlo()
139 rank, ShapeUtil::MakeShape(start_indices->shape().element_type(), {})); in MakeDynamicSliceHlo()
[all …]
Dhlo_creation_utils.h51 absl::Span<const int64> start_indices,
79 HloInstruction* operand, HloInstruction* start_indices,
87 HloInstruction* start_indices);
/external/tensorflow/tensorflow/lite/kernels/internal/
Dstrided_slice_logic.h48 p->start_indices[i + pad_count] = p->start_indices[i]; in StridedSlicePadIndices()
52 p->start_indices[i] = 0; in StridedSlicePadIndices()
77 const auto* start_indices = params.start_indices; in StartForAxis() local
80 int start = start_indices[axis]; in StartForAxis()
171 const std::vector<int>& start_indices, const std::vector<int>& stop_indices, in BuildStridedSliceParams() argument
174 const int dims_count = start_indices.size(); in BuildStridedSliceParams()
180 op_params.start_indices[i] = start_indices[i]; in BuildStridedSliceParams()
/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) =
Dfft_ops.cc211 const Eigen::DSizes<Eigen::DenseIndex, FFTRank + 1> start_indices; in DoFFT() local
215 full_fft.slice(start_indices, input_slice_sizes).device(device) = in DoFFT()
216 input.slice(start_indices, input_slice_sizes); in DoFFT()
224 full_fft.slice(start_indices, input_slice_sizes).device(device) = in DoFFT()
225 full_fft.slice(start_indices, input_slice_sizes) in DoFFT()
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dtensor_list_ops.cc224 std::vector<xla::XlaOp> start_indices(shape.dims(), in Compile() local
226 start_indices[0] = index; in Compile()
230 xla::XlaOp read = xla::DynamicSlice(buffer, start_indices, slice_shape); in Compile()
362 std::vector<xla::XlaOp> start_indices(elem_shape.dims() + 1, in Compile() local
364 start_indices[0] = index; in Compile()
372 start_indices), in Compile()
405 std::vector<xla::XlaOp> start_indices(elem_shape.dims() + 1, in Compile() local
407 start_indices[0] = index; in Compile()
416 BuildTensorList(xla::DynamicUpdateSlice(buffer, update, start_indices), in Compile()
450 std::vector<xla::XlaOp> start_indices(shape.dims(), in Compile() local
[all …]
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_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()
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()
Dreverse_sequence_op.cc96 xla::XlaOp start_indices = in Compile() local
122 xla::Gather(input, start_indices, dnums, slice_sizes)); in Compile()
Dfft_ops.cc85 std::vector<int64> start_indices(input_shape.dims(), 0); in Compile() local
87 input = xla::Pad(xla::Slice(input, start_indices, slice_sizes, strides), in Compile()
/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dsorting.cc41 std::vector<int64> start_indices(input_shape.dimensions_size(), 0); in TopK() local
46 XlaOp values = Neg(Slice(GetTupleElement(sort_result, 0), start_indices, in TopK()
48 XlaOp indices = Slice(GetTupleElement(sort_result, 1), start_indices, in TopK()
/external/tensorflow/tensorflow/lite/toco/
Dmodel.h1058 std::vector<int> start_indices;
1073 start_indices = other.start_indices;
1086 CHECK_GE(dim_count, start_indices.size());
1087 CHECK_EQ(start_indices.size(), stop_indices.size());
1090 for (int i = start_indices.size(); i < dim_count; i++) {
1091 start_indices.push_back(0);
1100 CHECK_EQ(start_indices.size(), stop_indices.size());
1103 std::reverse(start_indices.begin(), start_indices.end());
1108 (32 - start_indices.size());
1111 (32 - start_indices.size());
[all …]
/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/compiler/xla/service/gpu/
Dcudnn_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()
347 std::vector<int64> start_indices(new_backward_conv->shape().dimensions_size(), in CanonicalizeBackwardInputConvolution() local
361 start_indices[dim] += padding_low - padding_high; in CanonicalizeBackwardInputConvolution()
371 ShapeInference::InferSliceShape(new_backward_conv->shape(), start_indices, in CanonicalizeBackwardInputConvolution()
380 start_indices, limit_indices, strides)); in CanonicalizeBackwardInputConvolution()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.h346 XlaOp Slice(const XlaOp& operand, absl::Span<const int64> start_indices,
354 XlaOp DynamicSlice(const XlaOp& operand, const XlaOp& start_indices,
357 absl::Span<const XlaOp> start_indices,
362 const XlaOp& start_indices);
364 absl::Span<const XlaOp> start_indices);
543 XlaOp Gather(const XlaOp& input, const XlaOp& start_indices,
734 absl::Span<const int64> start_indices,
741 friend XlaOp DynamicSlice(const XlaOp& operand, const XlaOp& start_indices,
744 absl::Span<const XlaOp> start_indices,
748 const XlaOp& start_indices);
[all …]
Dxla_builder.cc722 absl::Span<const int64> start_indices, in Slice() argument
730 operand_shape, start_indices, limit_indices, strides)); in Slice()
732 for (int i = 0; i < start_indices.size(); i++) { in Slice()
734 slice_config->set_start(start_indices[i]); in Slice()
758 XlaOp XlaBuilder::DynamicSlice(const XlaOp& operand, const XlaOp& start_indices, in DynamicSlice() argument
765 GetShape(start_indices)); in DynamicSlice()
776 {operand, start_indices}); in DynamicSlice()
781 absl::Span<const XlaOp> start_indices, in DynamicSlice() argument
789 GetOperandShapes(start_indices)); in DynamicSlice()
803 operands.insert(operands.end(), start_indices.begin(), start_indices.end()); in DynamicSlice()
[all …]
/external/tensorflow/tensorflow/compiler/xla/
Dliteral.cc715 const Shape& result_shape, absl::Span<const int64> start_indices) const { in SliceInternal()
721 new_indices[i] = indices[i] + start_indices[i]; in SliceInternal()
729 Literal LiteralBase::Slice(absl::Span<const int64> start_indices, in Slice() argument
735 CHECK_GE(start_indices[dnum], 0); in Slice()
738 int64 dimension = limit_indices[dnum] - start_indices[dnum]; in Slice()
747 return SliceInternal<bool>(result_shape, start_indices); in Slice()
749 return SliceInternal<uint8>(result_shape, start_indices); in Slice()
751 return SliceInternal<uint16>(result_shape, start_indices); in Slice()
753 return SliceInternal<uint32>(result_shape, start_indices); in Slice()
755 return SliceInternal<uint64>(result_shape, start_indices); in Slice()
[all …]
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_client.py1317 def Slice(self, operand, start_indices, limit_indices, strides=None): argument
1333 start_indices = list(start_indices)
1334 strides = [1] * len(start_indices)
1335 return self._client.Slice(operand, start_indices, limit_indices, strides)
1353 def DynamicSlice(self, operand, start_indices, slice_sizes): argument
1366 return self._client.DynamicSlice(operand, start_indices, slice_sizes)
1368 def DynamicUpdateSlice(self, operand, update, start_indices): argument
1380 return self._client.DynamicUpdateSlice(operand, update, start_indices)
1775 def Gather(self, a, start_indices, dimension_numbers, slice_sizes): argument
1777 return self._client.Gather(a, start_indices, dimension_numbers, slice_sizes)
Dlocal_computation_builder.cc516 absl::Span<const int64> start_indices, in Slice() argument
519 return xla::Slice(operand.op(), start_indices, limit_indices, strides); in Slice()
529 const LocalOp& start_indices, in DynamicSlice() argument
531 return xla::DynamicSlice(operand.op(), start_indices.op(), slice_sizes); in DynamicSlice()
536 const LocalOp& start_indices) { in DynamicUpdateSlice() argument
537 return xla::DynamicUpdateSlice(operand.op(), update.op(), start_indices.op()); in DynamicUpdateSlice()
755 const LocalOp& input, const LocalOp& start_indices, in Gather() argument
758 return xla::Gather(input.op(), start_indices.op(), dimension_numbers, in Gather()
Dlocal_computation_builder.h262 LocalOp Slice(const LocalOp& operand, absl::Span<const int64> start_indices,
269 LocalOp DynamicSlice(const LocalOp& operand, const LocalOp& start_indices,
273 const LocalOp& start_indices);
373 LocalOp Gather(const LocalOp& input, const LocalOp& start_indices,

12