Home
last modified time | relevance | path

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

12

/external/tensorflow/tensorflow/compiler/mlir/hlo/tests/
Dhlo-legalize-gather-to-torch-index-select.mlir10 … index_vector_dim = 2 : i64, offset_dims = dense<2> : tensor<1xi64>, start_index_map = dense<0> : …
23 …x_vector_dim = 0 : i64, offset_dims = dense<[0, 1]> : tensor<2xi64>, start_index_map = dense<0> : …
32 … index_vector_dim = 2 : i64, offset_dims = dense<2> : tensor<1xi64>, start_index_map = dense<0> : …
39 … index_vector_dim = 2 : i64, offset_dims = dense<2> : tensor<1xi64>, start_index_map = dense<0> : …
Doptimize-hlo.mlir13 start_index_map = dense<0> : tensor<1xi64>
34 start_index_map = dense<0> : tensor<1xi64>
57 start_index_map = dense<0> : tensor<1xi64>
Dcanonicalize.mlir1022 start_index_map = dense<[0, 2]> : tensor<2xi64>},
1037 start_index_map = dense<[2]> : tensor<1xi64>},
1052 start_index_map = dense<[0, 2]> : tensor<2xi64>},
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/
Doptimize_mhlo.cc74 if (!dimension_numbers.start_index_map() || in matchAndRewrite()
75 dimension_numbers.start_index_map().getType().getRank() != 1 || in matchAndRewrite()
76 dimension_numbers.start_index_map().getType().getDimSize(0) != 1 || in matchAndRewrite()
77 dimension_numbers.start_index_map() in matchAndRewrite()
Dlegalize_gather_to_torch_index_select.cc63 if (dimension_numbers.start_index_map().getType().getRank() != 1 || in matchAndRewrite()
64 dimension_numbers.start_index_map() in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dattribute_importer.cc54 std::vector<int64_t> start_index_map(dnums.start_index_map().begin(), in ConvertGatherDimensionNumbers() local
55 dnums.start_index_map().end()); in ConvertGatherDimensionNumbers()
58 Convert(start_index_map, builder), in ConvertGatherDimensionNumbers()
Dhlo_utils.cc271 auto start_index_map = CreateDenseIntElementsAttrFromVector( in CreateGatherDimensionNumbers() local
272 llvm::SmallVector<int64, 4>{input.start_index_map().begin(), in CreateGatherDimensionNumbers()
273 input.start_index_map().end()}, in CreateGatherDimensionNumbers()
280 offset_dims, collapsed_slice_dims, start_index_map, index_vector_dim, in CreateGatherDimensionNumbers()
Dmlir_hlo_to_hlo.cc307 auto start_index_map = ConvertDenseIntAttr(input.start_index_map()); in Convert_dimension_numbers() local
308 std::copy(start_index_map.begin(), start_index_map.end(), in Convert_dimension_numbers()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dgather_scatter_handler.cc221 std::vector<int64> start_index_map(dnums.start_index_map().begin(), in ParititonTrivialIndexedOperandDimension() local
222 dnums.start_index_map().end()); in ParititonTrivialIndexedOperandDimension()
224 operand, start_index_map, gather->gather_slice_sizes()) && in ParititonTrivialIndexedOperandDimension()
237 operand, indices, operand.state().partition_id, start_index_map, in ParititonTrivialIndexedOperandDimension()
385 operand_offsets[dnums.start_index_map(start_idx)])) in PartitionIndexParallelDimensions()
386 : operand_offsets[dnums.start_index_map(start_idx)]; in PartitionIndexParallelDimensions()
/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_sharding_util.cc952 std::vector<int64> start_index_map(dnums.start_index_map().begin(), in GatherOutputShardingFromDataOperand() local
953 dnums.start_index_map().end()); in GatherOutputShardingFromDataOperand()
958 start_index_map, offset_dims, hlo.gather_slice_sizes()); in GatherOutputShardingFromDataOperand()
966 std::vector<int64> start_index_map(dnums.start_index_map().begin(), in GatherDataOperandShardingFromOutput() local
967 dnums.start_index_map().end()); in GatherDataOperandShardingFromOutput()
986 collapsed_slice_dims, start_index_map, offset_dims, in GatherDataOperandShardingFromOutput()
1322 if (hlo.gather_slice_sizes()[dnums.start_index_map(i)] == 1) { in GetGatherBatchParallelDims()
1324 operand_parallel_dims.push_back(dnums.start_index_map(i)); in GetGatherBatchParallelDims()
1377 const int64 operand_parallel_dim = dnums.start_index_map(i); in GatherOutputAlignedOperandParallelDims()
Dgather_expander.cc132 int64 index_vector_dim_index = FindIndex(dim_numbers.start_index_map(), i); in ExpandIndexVectorIntoOperandSpace()
Dindexed_array_analysis.cc267 dim_numbers.collapsed_slice_dims(0) != dim_numbers.start_index_map(0)) { in ComputeArrayForGather()
290 int64 source_dim = dim_numbers.start_index_map(0); in ComputeArrayForGather()
Dhlo_instructions.cc2744 string start_index_map = in GatherDimensionNumbersToString() local
2746 StrJoin(gather_dimension_numbers.start_index_map(), ","), "}"); in GatherDimensionNumbersToString()
2751 {offset_dims, collapsed_slice_dims, start_index_map, index_vector_dim}, in GatherDimensionNumbersToString()
2758 absl::Span<const int64> start_index_map, int64 index_vector_dim) { in MakeGatherDimNumbers() argument
2766 for (int64 gather_dim_to_input_dim : start_index_map) { in MakeGatherDimNumbers()
Dhlo_parser_test.cc914 …} %start_indices), offset_dims={4,5,6,7,8}, collapsed_slice_dims={}, start_index_map={0,1,2,3,4}, … in CreateTestCases()
926 …} %start_indices), offset_dims={4,5,6,7,8}, collapsed_slice_dims={}, start_index_map={0,1,2,3,4}, … in CreateTestCases()
1516 …r, start_indices), offset_dims={4,5,6,7,8}, collapsed_slice_dims={}, start_index_map={0,1,2,3,4}, … in CreateTestCases()
Dshape_inference.cc3253 int64 operand_dim_for_start_index_i = dim_numbers.start_index_map(i); in ValidateGatherDimensionNumbers()
3263 dim_numbers.start_index_map().begin(), in ValidateGatherDimensionNumbers()
3264 dim_numbers.start_index_map().end()); in ValidateGatherDimensionNumbers()
3273 StrJoin(dim_numbers.start_index_map(), ", ")); in ValidateGatherDimensionNumbers()
/external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/
Dhlo_ops_base_structs.td58 StructFieldAttr<"start_index_map", I64ElementsAttr>,
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dlegalize_hlo.cc1038 auto start_index_map = gather_op.dimension_numbers().start_index_map(); in matchAndRewrite() local
1041 if (!IsIotaAttr(start_index_map, start_indices_type.getShape().back()) || in matchAndRewrite()
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_client.py647 self.start_index_map = []
/external/tensorflow/tensorflow/compiler/tests/
Dxla_ops_test.py95 dimension_numbers.start_index_map.extend([0])
/external/tensorflow/tensorflow/compiler/xla/tests/
Dtest_utils.cc584 ? use->gather_dimension_numbers().start_index_map() in CreateLiteralForConstrainedUses()
/external/tensorflow/tensorflow/compiler/xla/g3doc/
Doperation_semantics.md1431 <b> `gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map
1455 | `start_index_map` | `ArraySlice<int64>` | A map that describes how to |
1504 scattering `S` using `start_index_map`. More precisely:
1506 1. `S`<sub>`in`</sub>[`start_index_map`[`k`]] = `S`[`k`] if `k` <
1507 `start_index_map.size`.
1530 are sorted (in ascending `start_index_map` order) by the user. If they are not
1546 - We use `start_index_map` to map the starting index (whose size may be less
1630 these are expanded in accordance to "Gather Index Mapping" (`start_index_map` in
2493 are sorted (in ascending `start_index_map` order) by the user. If they are not
/external/tensorflow/tensorflow/compiler/xla/
Dxla_data.proto490 // This is interpreted as a map from i to start_index_map[i]. It
493 repeated int64 start_index_map = 3; field
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/hlo_to_lhlo_with_xla/
Dops.mlir820 // CHECK-SAME: start_index_map = dense<[0, 1]> : tensor<2xi64>},
824 … index_vector_dim = 1 : i64, offset_dims = dense<1> : tensor<1xi64>, start_index_map = dense<[0, 1…
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/translate/
Dexport.mlir489 // CHECK-SAME: start_index_map={0,1}
493 … index_vector_dim = 1 : i64, offset_dims = dense<1> : tensor<1xi64>, start_index_map = dense<[0, 1…
Dimport.hlotxt420 // CHECK-SAME: start_index_map = dense<[0, 1]> : tensor<2xi64>
427 start_index_map={0,1},

12