Searched refs:target_dims (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.cc | 1797 absl::Span<const int64> target_dims, in CreateMatchingShardingOnDims() argument
1799 CHECK(target_dims.size() == source_dims.size()) in CreateMatchingShardingOnDims()
1806 for (int i = 0, end = target_dims.size(); i < end; ++i) { in CreateMatchingShardingOnDims()
1808 tile_dims[target_dims[i]] = in CreateMatchingShardingOnDims()
1824 target_dims, source_sharding, source_dims); in CreateMatchingShardingOnDims()
1827 target_dims, source_sharding, source_dims); in CreateMatchingShardingOnDims()
1872 auto target_dims = indices_parallel_dims_ordered_as_operand; in GatherOperandsShardedAcrossParallelDims() local
1876 std::swap(to_adjust_dims, target_dims); in GatherOperandsShardedAcrossParallelDims()
1884 int64 target_dim = target->tile_assignment().dim(target_dims[i]); in GatherOperandsShardedAcrossParallelDims()
1911 to_adjust_dims, *target, target_dims); in GatherOperandsShardedAcrossParallelDims()
[all …]
|
| D | spmd_partitioner_util.h | 410 absl::Span<const int64> target_dims,
|