Home
last modified time | relevance | path

Searched refs:limit_indices (Results 1 – 25 of 36) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/client/lib/
Dsorting.cc64 std::vector<int64> limit_indices(input_dims.begin(), input_dims.end()); in TopK() local
65 limit_indices[last_dim] = k; in TopK()
69 limit_indices, strides); in TopK()
71 limit_indices, strides); in TopK()
127 std::vector<int64> limit_indices(input_dims.begin(), input_dims.end()); in TopKWithPartitions() local
131 limit_indices[last_dim] = k; in TopKWithPartitions()
133 limit_indices, strides); in TopKWithPartitions()
135 limit_indices, strides); in TopKWithPartitions()
142 std::vector<int64> limit_indices(input_dims.begin(), input_dims.end()); in TopKWithPartitions() local
145 limit_indices[last_dim] = per_partition_size; in TopKWithPartitions()
[all …]
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dunpack_op.cc68 std::vector<int64> limit_indices(input_shape.dims()); in Compile() local
71 limit_indices[i] = input_shape.dim_size(i); in Compile()
76 limit_indices[axis] = i + 1; in Compile()
77 auto slice = xla::Slice(input, start_indices, limit_indices, strides); in Compile()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_conv_padding_legalization.cc95 std::vector<int64> limit_indices(input->shape().dimensions().begin(), in MaybePaddedAndSlicedInput() local
107 limit_indices[dim] -= -conv_window->dimensions(i).padding_high(); in MaybePaddedAndSlicedInput()
113 MakeSliceHlo(input, start_indices, limit_indices, strides).ValueOrDie(); in MaybePaddedAndSlicedInput()
353 std::vector<int64> limit_indices( in CanonicalizeBackwardInputConvolution() local
373 limit_indices[dim] -= padding_high - padding_low; in CanonicalizeBackwardInputConvolution()
380 limit_indices, strides) in CanonicalizeBackwardInputConvolution()
388 start_indices, limit_indices, strides)); in CanonicalizeBackwardInputConvolution()
/external/tensorflow/tensorflow/compiler/mlir/hlo/tests/
Doptimize-hlo.mlir46 …// CHECK-DAG: [[SLICE1:%.+]] = "mhlo.slice"(%arg1) {limit_indices = dense<1> : tensor<1xi64>, star…
48 …// CHECK-DAG: [[SLICE2:%.+]] = "mhlo.slice"(%arg1) {limit_indices = dense<2> : tensor<1xi64>, star…
Dcanonicalize.mlir244 // CHECK-DAG-SAME: limit_indices = dense<3> : tensor<1xi64>
256 // CHECK-DAG-SAME: limit_indices = dense<[2, 4]> : tensor<2xi64>
269 …%0 = "mhlo.slice"(%arg0) { limit_indices = dense<[2, 2]> : tensor<2xi64>, start_indices = dense<[0…
279 …%1 = "mhlo.slice"(%0) { limit_indices = dense<[3]> : tensor<1xi64>, start_indices = dense<[1]> : t…
287 …%1 = "mhlo.slice"(%0) { limit_indices = dense<[3]> : tensor<1xi64>, start_indices = dense<[1]> : t…
295 …%1 = "mhlo.slice"(%0) { limit_indices = dense<[4]> : tensor<1xi64>, start_indices = dense<[1]> : t…
306 …%1 = "mhlo.slice"(%0) { limit_indices = dense<[3, 4]> : tensor<2xi64>, start_indices = dense<[1, 2…
316 …%1 = "mhlo.slice"(%0) { limit_indices = dense<[1, 4]> : tensor<2xi64>, start_indices = dense<[0, 0…
326 …%1 = "mhlo.slice"(%0) { limit_indices = dense<[4, 3]> : tensor<2xi64>, start_indices = dense<[0, 2…
334 …%1 = "mhlo.slice"(%0) { limit_indices = dense<[0]> : tensor<1xi64>, start_indices = dense<[0]> : t…
[all …]
Dops.mlir719 …%0 = "mhlo.slice"(%arg0) {start_indices = dense<[1, 0]> : tensor<2xi64>, limit_indices = dense<[2,…
726 …// expected-error@+1 {{failed to verify that all of {start_indices, limit_indices, strides} have s…
727 …%0 = "mhlo.slice"(%arg0) {start_indices = dense<[1, 2]> : tensor<2xi64>, limit_indices = dense<[2,…
735 …%0 = "mhlo.slice"(%arg0) {start_indices = dense<[1, 0]> : tensor<2xi64>, limit_indices = dense<[2,…
745 limit_indices = dense<[[2, 4]]> : tensor<1x2xi64>,
757 limit_indices = dense<[2, 4, 0]> : tensor<3xi64>,
/external/tensorflow/tensorflow/compiler/xla/
Darray.h481 std::vector<int64> limit_indices;
483 from.dimensions().begin(), std::back_inserter(limit_indices),
488 if (array_impl::all_inside_range(index, start_indices, limit_indices)) {
Dliteral.cc913 absl::Span<const int64> limit_indices) const { in Slice()
919 CHECK_LE(limit_indices[dnum], shape().dimensions(dnum)) in Slice()
921 int64 dimension = limit_indices[dnum] - start_indices[dnum]; in Slice()
Dliteral.h335 absl::Span<const int64> limit_indices) const;
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/
Dlegalize-tf.mlir1889 …// CHECK: %[[SLICED:.*]] = "mhlo.slice"(%arg0) {limit_indices = dense<[2, 8]> : tensor<2xi64>, sta…
1898 …// CHECK: %[[SLICED:.*]] = "mhlo.slice"(%arg0) {limit_indices = dense<5> : tensor<1xi64>, start_in…
2608 // CHECK-DAG-SAME: {limit_indices = dense<1> : tensor<1xi64>,
2630 // CHECK-DAG-SAME: {limit_indices = dense<1> : tensor<1xi64>,
2647 // CHECK-DAG-SAME: {limit_indices = dense<1> : tensor<1xi64>,
2665 // CHECK-DAG-SAME: {limit_indices = dense<1> : tensor<1xi64>,
2672 // CHECK-DAG-SAME: {limit_indices = dense<2> : tensor<1xi64>,
2693 // CHECK-DAG-SAME: {limit_indices = dense<1> : tensor<1xi64>,
2697 …// CHECK: %[[SLICED_START2:.*]] = "mhlo.slice"(%[[START_I64]]) {limit_indices = dense<2> : tensor<…
2735 // CHECK-DAG-SAME: limit_indices = dense<[3, 7]>
[all …]
Dlegalize-tf-with-tf2xla.mlir175 // CHECK-DAG-SAME: limit_indices = dense<1> : tensor<1xi64>
182 // CHECK-DAG-SAME: limit_indices = dense<2> : tensor<1xi64>
/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_creation_utils.cc83 absl::Span<const int64> limit_indices, in MakeSliceHlo() argument
88 limit_indices, strides)); in MakeSliceHlo()
90 slice_shape, operand, start_indices, limit_indices, strides)); in MakeSliceHlo()
Dhlo_creation_utils.h57 absl::Span<const int64> limit_indices,
Dhlo_evaluator.cc2190 std::vector<int64> limit_indices(indices.begin(), indices.end()); in HandleSort() local
2191 absl::c_for_each(limit_indices, [](int64& index) { ++index; }); in HandleSort()
2192 limit_indices[sort_dim] = sort_dim_elements; in HandleSort()
2198 .Slice(indices, limit_indices) in HandleSort()
/external/tensorflow/tensorflow/compiler/mlir/xla/ir/
Dmlir_hlo_builder.h243 absl::Span<const int64> limit_indices,
Dmlir_hlo_builder.cc510 absl::Span<const int64> limit_indices, absl::Span<const int64> strides) { in SliceInternal() argument
513 GetI64ElementsAttr(limit_indices, &builder_), in SliceInternal()
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dlegalize_hlo.cc281 auto limit_indices = slice_op.limit_indices(); in matchAndRewrite() local
284 limit_indices.getValues<APInt>())) { in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/
Dlhlo_ops.td302 [AllTypesMatch<["start_indices", "limit_indices", "strides"]>]> {
307 I64ElementsAttr:$limit_indices,
Dhlo_ops.td707 AllTypesMatch<["start_indices", "limit_indices", "strides"]>,
712 I64ElementsAttr:$limit_indices,
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/IR/
Dhlo_ops.cc2394 slice.limit_indices().getType() != attr_ty || in inferReturnTypes()
2407 SmallVector<int64_t, 4> limit(slice.limit_indices().getValues<int64_t>()); in inferReturnTypes()
2450 auto limit = llvm::to_vector<6>(op->limit_indices().getValues<int64_t>()); in FoldSlice()
2535 auto limit = slice.limit_indices().getIntValues(); in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.h485 absl::Span<const int64> limit_indices,
489 absl::Span<const int64> limit_indices,
1091 absl::Span<const int64> limit_indices,
1746 absl::Span<const int64> limit_indices,
Dxla_builder.cc953 absl::Span<const int64> limit_indices, in Slice() argument
959 limit_indices, strides)); in Slice()
960 return SliceInternal(shape, operand, start_indices, limit_indices, strides); in Slice()
966 absl::Span<const int64> limit_indices, in SliceInternal() argument
973 slice_config->set_limit(limit_indices[i]); in SliceInternal()
4155 absl::Span<const int64> limit_indices, in Slice() argument
4157 return operand.builder()->Slice(operand, start_indices, limit_indices, in Slice()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc1266 std::vector<int64> limit_indices(hlo_dims.begin(), hlo_dims.end()); in SliceFirstK() local
1268 limit_indices[slice_dim] = k; in SliceFirstK()
1272 output_shape, hlo, start_indices, limit_indices, strides)); in SliceFirstK()
Dspmd_partitioner.cc1612 std::vector<int64> limit_indices = hlo->slice_limits(); in HandleSlice() local
1618 limit_indices[i] = operand_shape.dimensions(i) + dim.padding_high(); in HandleSlice()
1620 limit_indices[i] != operand_shape.dimensions(i)) { in HandleSlice()
1630 limit_indices, strides)); in HandleSlice()
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/hlo_to_lhlo_with_xla/
Dops.mlir805 // CHECK-SAME: limit_indices = dense<[2, 4]> : tensor<2xi64>,
810 …%0 = "mhlo.slice"(%arg) {start_indices = dense<[1, 0]> : tensor<2xi64>, limit_indices = dense<[2, …

12