Home
last modified time | relevance | path

Searched refs:group_dims (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/core/kernels/
Dedit_distance_op.cc148 std::vector<int64> group_dims(truth_st_shape.dims() - 1); in Compute() local
149 std::iota(group_dims.begin(), group_dims.end(), 0); in Compute()
152 for (int d = 0; d < static_cast<int>(group_dims.size()); ++d) { in Compute()
168 auto hypothesis_grouper = hypothesis.group(group_dims); in Compute()
169 auto truth_grouper = truth.group(group_dims); in Compute()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.h299 std::vector<int64> group_dims, in GroupedSharding()
303 group_dims(std::move(group_dims)), in GroupedSharding()
308 std::vector<int64> group_dims; member
316 absl::Span<const int64> group_dims,
321 absl::Span<const int64> group_dims);
355 absl::Span<const int64> group_dims, absl::Span<const int64> group_dim_sizes,
Dspmd_partitioner_util.cc1379 absl::Span<const int64> group_dims) { in GroupShardingOnDims() argument
1380 std::vector<int64> group_dim_shards(group_dims.size(), 1); in GroupShardingOnDims()
1381 return GroupShardingOnDims(sharding, group_dims, group_dim_shards); in GroupShardingOnDims()
1385 absl::Span<const int64> group_dims, in GroupShardingOnDims() argument
1390 std::vector<int64> group_dim_sizes(group_dims.size()); in GroupShardingOnDims()
1391 for (int64 i = 0; i < group_dims.size(); ++i) { in GroupShardingOnDims()
1392 CHECK_EQ(grouped_tiling_dims[group_dims[i]] % group_dim_shards[i], 0); in GroupShardingOnDims()
1394 grouped_tiling_dims[group_dims[i]] / group_dim_shards[i]; in GroupShardingOnDims()
1395 grouped_tiling_dims[group_dims[i]] = group_dim_shards[i]; in GroupShardingOnDims()
1402 for (int64 i = 0; i < group_dims.size(); ++i) { in GroupShardingOnDims()
[all …]
Ddot_handler.cc2043 if (other_grouped && other_grouped->group_dims.size() == 1 && in PartitionDotGroupOnNonContracting()
2044 other_grouped->group_dims[0] == other.base_shape().rank()) { in PartitionDotGroupOnNonContracting()
2055 other.sharding(), other_grouped->group_dims)) in PartitionDotGroupOnNonContracting()
2190 std::vector<int64> group_dims; in PartitionDotGroupOnContracting() local
2193 group_dims = std::move(*found_dims); in PartitionDotGroupOnContracting()
2199 group_dims.push_back(dim.output); in PartitionDotGroupOnContracting()
2203 group_dims.push_back(dim.output); in PartitionDotGroupOnContracting()
2207 group_dims.push_back(dim.output); in PartitionDotGroupOnContracting()
2211 if (!group_dims.empty()) { in PartitionDotGroupOnContracting()
2213 GroupShardingOnDims(output_sharding, group_dims), lhs_grouped); in PartitionDotGroupOnContracting()
[all …]
Dspmd_partitioner.cc456 std::vector<int64> group_dims(target.tile_assignment().num_dimensions() - in ReshardNoCache() local
458 std::iota(group_dims.begin(), group_dims.end(), 0); in ReshardNoCache()
459 auto target_grouped = GroupShardingOnDims(target, group_dims); in ReshardNoCache()
461 hlo_, state_.partition_id, target_grouped.device_groups, group_dims, in ReshardNoCache()
3083 std::vector<int64> group_dims( in HandleRng() local
3085 std::iota(group_dims.begin(), group_dims.end(), 0); in HandleRng()
3086 auto sharding_grouped = GroupShardingOnDims(hlo->sharding(), group_dims); in HandleRng()
/external/tensorflow/tensorflow/core/util/sparse/
Dgroup_iterator.h81 GroupIterable(Tensor ix, Tensor vals, int dims, const VarDimArray& group_dims) in GroupIterable() argument
86 group_dims_(group_dims.begin(), group_dims.end()) {} in GroupIterable()