| /external/tensorflow/tensorflow/core/kernels/ |
| D | split_lib_cpu.cc | 32 const Eigen::DSizes<Eigen::DenseIndex, NDims>& slice_sizes) { in operator ()() argument
34 output = input.slice(slice_indices, slice_sizes); in operator ()()
36 output.device(d) = input.slice(slice_indices, slice_sizes); in operator ()()
53 const Eigen::DSizes<Eigen::DenseIndex, NDims>& slice_sizes) { in DEFINE_CPU_KERNELS()
54 output.device(d) = input.slice(slice_indices, slice_sizes); in DEFINE_CPU_KERNELS()
|
| D | split_lib.h | 31 const Eigen::DSizes<Eigen::DenseIndex, 2>& slice_sizes);
39 const Eigen::DSizes<Eigen::DenseIndex, NDims>& slice_sizes);
48 const Eigen::DSizes<Eigen::DenseIndex, NDims>& slice_sizes);
58 const Eigen::DSizes<Eigen::DenseIndex, NDims>& slice_sizes);
|
| D | slice_op.h | 32 const Eigen::DSizes<Eigen::DenseIndex, NDIMS>& slice_sizes) { in operator()
42 sizes[i] = slice_sizes[i]; in operator()
46 output.device(d) = input.slice(slice_indices, slice_sizes); in operator()
|
| D | split_op.cc | 169 Eigen::DSizes<Eigen::DenseIndex, NDims> slice_sizes; in operator ()() local
173 slice_sizes[j] = sizes[j]; in operator ()()
180 result_shaped = input_reshaped.slice(slice_indices, slice_sizes); in operator ()()
185 input_reshaped, slice_indices, slice_sizes); in operator ()()
379 Eigen::DSizes<Eigen::DenseIndex, 3> slice_sizes; in Compute() local
382 slice_sizes[j] = sizes[j]; in Compute()
390 slice_indices, slice_sizes); in Compute()
|
| D | split_lib_gpu.cu.cc | 38 const Eigen::DSizes<Eigen::DenseIndex, NDims>& slice_sizes) { in operator ()() argument
39 To32Bit(output).device(d) = To32Bit(input).slice(slice_indices, slice_sizes); in operator ()()
47 const Eigen::DSizes<Eigen::DenseIndex, 2>& slice_sizes) { in operator ()() argument
48 To32Bit(output).device(d) = To32Bit(input).slice(slice_indices, slice_sizes); in operator ()()
|
| D | topk_op_gpu.h | 524 const Eigen::DSizes<Eigen::DenseIndex, 2> slice_sizes{num_rows, k};
526 To32Bit(temp_indices.matrix<int32>()).slice(slice_indices, slice_sizes);
528 To32Bit(temp_values.matrix<T>()).slice(slice_indices, slice_sizes);
|
| D | batch_kernels.cc | 162 Eigen::DSizes<Eigen::DenseIndex, 2> slice_sizes{size, suffix_dim_size}; in SplitCPU() local
165 slice_indices, slice_sizes); in SplitCPU()
360 const std::vector<int64> slice_sizes = {1}; in ConcatInputTensors() local
368 SplitCPU<type>(context, padding_source, slice_sizes, &slices); \ in ConcatInputTensors()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | reverse_sequence_op.cc | 117 auto slice_sizes = input_shape.dim_sizes(); in Compile() local
118 slice_sizes[batch_dim_] = 1; in Compile()
119 slice_sizes[seq_dim_] = 1; in Compile()
122 xla::Gather(input, start_indices, dnums, slice_sizes)); in Compile()
|
| D | fft_ops.cc | 61 absl::InlinedVector<int64, 4> slice_sizes = input_shape.dim_sizes(); in Compile() local
62 std::vector<std::pair<int64, int64>> padding_sizes(slice_sizes.size()); in Compile()
78 slice_sizes[index] = expected_sizes[i]; in Compile()
87 input = xla::Pad(xla::Slice(input, start_indices, slice_sizes, strides), in Compile()
|
| D | gather_op.cc | 117 std::vector<int64> slice_sizes; in XlaGather() local
118 slice_sizes.reserve(input_shape.dims()); in XlaGather()
128 slice_sizes.push_back(window_bound); in XlaGather()
145 *gather_output = xla::Gather(input, indices, dim_numbers, slice_sizes); in XlaGather()
|
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | slicing.cc | 153 std::vector<int64> slice_sizes(input_shape.rank(), 1); in TorchGather() local
160 return Gather(input, gather_indices, gather_dnums, slice_sizes); in TorchGather()
169 std::vector<int64> slice_sizes = input_shape.dimensions(); in TorchIndexSelect() local
170 slice_sizes[dim] = 1; in TorchIndexSelect()
180 return Gather(input, index, gather_dnums, slice_sizes); in TorchIndexSelect()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | gather_expander.cc | 234 absl::Span<const int64> slice_sizes, int64 gather_loop_trip_count, in CreateGatherLoopAccumulatorInitValue() argument
237 accumulator_state_shape_dims.reserve(1 + slice_sizes.size()); in CreateGatherLoopAccumulatorInitValue()
239 for (int64 i = 0; i < slice_sizes.size(); i++) { in CreateGatherLoopAccumulatorInitValue()
241 accumulator_state_shape_dims.push_back(slice_sizes[i]); in CreateGatherLoopAccumulatorInitValue()
|
| D | triangular_solve_expander.cc | 73 std::vector<int64> slice_sizes(ndims); in DiagonalBlocks() local
78 slice_sizes[i] = ShapeUtil::GetDimension(shape, i); in DiagonalBlocks()
80 slice_sizes[ndims - 2] = slice_sizes[ndims - 1] = block_size; in DiagonalBlocks()
86 diag_blocks = Gather(a, start_indices, dim_numbers, slice_sizes); in DiagonalBlocks()
|
| D | shape_inference.cc | 2267 absl::Span<const int64> slice_sizes, bool allow_scalar_indices) { in InferDynamicSliceShape() argument
2284 StrJoin(slice_sizes, ", ")); in InferDynamicSliceShape()
2311 StrJoin(slice_sizes, ", ")); in InferDynamicSliceShape()
2343 if (slice_sizes.size() != operand_shape.rank()) { in InferDynamicSliceShape()
2346 slice_sizes.size(), operand_shape.rank()); in InferDynamicSliceShape()
2349 for (int64 dim = 0; dim < slice_sizes.size(); ++dim) { in InferDynamicSliceShape()
2351 const int64 slice_dim_size = slice_sizes[dim]; in InferDynamicSliceShape()
2364 return ShapeUtil::MakeShape(operand_shape.element_type(), slice_sizes); in InferDynamicSliceShape()
2960 absl::Span<const int64> slice_sizes) { in InferGatherShape() argument
2999 if (slice_sizes.size() != input_shape.dimensions_size()) { in InferGatherShape()
[all …]
|
| D | shape_inference.h | 187 absl::Span<const int64> slice_sizes, bool allow_scalar_indices = true);
285 absl::Span<const int64> slice_sizes);
|
| D | hlo_creation_utils.cc | 123 absl::Span<const int64> slice_sizes) { in MakeDynamicSliceHlo() argument
143 operand->shape(), scalar_start_indices_shapes, slice_sizes)); in MakeDynamicSliceHlo()
145 dynamic_slice_shape, operand, scalar_start_indices, slice_sizes)); in MakeDynamicSliceHlo()
|
| D | hlo_creation_utils.h | 80 absl::Span<const int64> slice_sizes);
|
| D | hlo_instructions.cc | 2235 absl::Span<const int64> slice_sizes) in HloDynamicSliceInstruction() argument
2237 dynamic_slice_sizes_(slice_sizes.begin(), slice_sizes.end()) { in HloDynamicSliceInstruction()
2245 absl::Span<const int64> slice_sizes) in HloDynamicSliceInstruction() argument
2247 dynamic_slice_sizes_(slice_sizes.begin(), slice_sizes.end()) { in HloDynamicSliceInstruction()
2312 absl::Span<const int64> slice_sizes) in HloGatherInstruction() argument
2318 absl::c_copy(slice_sizes, std::back_inserter(gather_slice_sizes_)); in HloGatherInstruction()
|
| D | indexed_array_analysis.cc | 255 absl::Span<const int64> slice_sizes, Array* source, Array* indices) { in ComputeArrayForGather() argument
280 source->shape().dimensions(i) != slice_sizes[i]) { in ComputeArrayForGather()
283 << source->shape().dimensions(i) << " vs. " << slice_sizes[i] in ComputeArrayForGather()
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | dynamic_ops_test.cc | 119 const std::vector<int64>& slice_sizes, in RunR1() argument
142 DynamicSlice(input, absl::Span<const XlaOp>({starts}), slice_sizes); in RunR1()
150 const std::vector<int64>& slice_sizes, in RunR2() argument
172 DynamicSlice(input, starts, slice_sizes); in RunR2()
185 const std::vector<int64>& slice_sizes, in RunR3() argument
206 DynamicSlice(input, starts, slice_sizes); in RunR3()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.h | 355 absl::Span<const int64> slice_sizes);
358 absl::Span<const int64> slice_sizes);
545 absl::Span<const int64> slice_sizes);
742 absl::Span<const int64> slice_sizes);
745 absl::Span<const int64> slice_sizes);
970 absl::Span<const int64> slice_sizes);
1215 absl::Span<const int64> slice_sizes);
1219 absl::Span<const int64> slice_sizes);
1817 absl::Span<const int64> slice_sizes);
|
| D | xla_builder.cc | 759 absl::Span<const int64> slice_sizes) { in DynamicSlice() argument
768 operand_shape, {start_indices_shape}, slice_sizes)); in DynamicSlice()
771 for (int64 size : slice_sizes) { in DynamicSlice()
782 absl::Span<const int64> slice_sizes) { in DynamicSlice() argument
795 operand_shape, start_indices_shapes, slice_sizes)); in DynamicSlice()
798 for (int64 size : slice_sizes) { in DynamicSlice()
1842 absl::Span<const int64> slice_sizes) { in Gather() argument
1851 dimension_numbers, slice_sizes)); in Gather()
1855 for (int64 bound : slice_sizes) { in Gather()
2882 absl::Span<const int64> slice_sizes) { in DynamicSlice() argument
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | local_computation_builder.cc | 530 absl::Span<const int64> slice_sizes) { in DynamicSlice() argument
531 return xla::DynamicSlice(operand.op(), start_indices.op(), slice_sizes); in DynamicSlice()
757 absl::Span<const int64> slice_sizes) { in Gather() argument
759 slice_sizes); in Gather()
|
| D | local_computation_builder.h | 270 absl::Span<const int64> slice_sizes);
375 absl::Span<const int64> slice_sizes);
|
| D | xla_client.py | 1353 def DynamicSlice(self, operand, start_indices, slice_sizes): argument
1366 return self._client.DynamicSlice(operand, start_indices, slice_sizes)
1775 def Gather(self, a, start_indices, dimension_numbers, slice_sizes): argument
1777 return self._client.Gather(a, start_indices, dimension_numbers, slice_sizes)
|