Searched refs:indices_parallel_dims (Results 1 – 5 of 5) sorted by relevance
1306 absl::InlinedVector<int64, 1> indices_parallel_dims; in GetGatherBatchParallelDims() local1318 if (absl::c_linear_search(indices_parallel_dims, index_parallel_dim)) { in GetGatherBatchParallelDims()1323 indices_parallel_dims.push_back(index_parallel_dim); in GetGatherBatchParallelDims()1329 absl::c_sort(indices_parallel_dims); in GetGatherBatchParallelDims()1330 if (!indices_parallel_dims.empty()) { in GetGatherBatchParallelDims()1331 return GatherParallelDims{indices_parallel_dims, operand_parallel_dims, in GetGatherBatchParallelDims()1340 auto indices_parallel_dims = parallel_dim.indices_parallel_dims; in GatherParallelOutputDims() local1349 if (absl::c_binary_search(indices_parallel_dims, index_dim)) { in GatherParallelOutputDims()1362 CHECK_LE(parallel_dims.indices_parallel_dims.size(), in GatherOutputAlignedOperandParallelDims()1375 for (int j = 0; j < parallel_dims.indices_parallel_dims.size(); ++j) { in GatherOutputAlignedOperandParallelDims()[all …]
34 absl::InlinedVector<int64, 1> indices_parallel_dims; member
469 absl::MakeConstSpan(parallel_dims.indices_parallel_dims), in InferGatherParallelShardingFromOperands()
335 auto indices_parallel_dims = parallel_dims->indices_parallel_dims; in PartitionIndexParallelDimensions() local342 GroupShardingOnDims(indices_sharding, indices_parallel_dims); in PartitionIndexParallelDimensions()352 int indices_idx = indices_parallel_dims[i]; in PartitionIndexParallelDimensions()428 indices_sharding.NumTiles(indices_parallel_dims)) { in PartitionIndexParallelDimensions()
1835 auto& indices_parallel_dims = parallel_dims.indices_parallel_dims; in GatherOperandsShardedAcrossParallelDims() local1837 if (indices_parallel_dims.size() != operand_parallel_dims.size()) { in GatherOperandsShardedAcrossParallelDims()1842 int idx_parallel_tiles_num = new_index_shard.NumTiles(indices_parallel_dims); in GatherOperandsShardedAcrossParallelDims()