Searched refs:selected_dims (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.h | 212 int64* next_channel_id, absl::Span<const int64> selected_dims,
219 int64* next_channel_id, absl::Span<const int64> selected_dims,
234 int64* next_channel_id, absl::Span<const int64> selected_dims,
238 int64* next_channel_id, absl::Span<const int64> selected_dims,
|
| D | spmd_partitioner.cc | 3501 int64* next_channel_id, absl::Span<const int64> selected_dims, in AllGatherShards() argument
3504 selected_dims, collectives_creator, in AllGatherShards()
3510 int64* next_channel_id, absl::Span<const int64> selected_dims, in AllGatherShardsInternal() argument
3512 if (selected_dims.empty()) { in AllGatherShardsInternal()
3526 for (auto it = selected_dims.rbegin(); it != selected_dims.rend(); ++it) { in AllGatherShardsInternal()
3541 GetPartitionGroupsForReplication(sharding, selected_dims); in AllGatherShardsInternal()
3552 absl::c_linear_search(selected_dims, i)) { in AllGatherShardsInternal()
3575 !absl::c_linear_search(selected_dims, i - split_dims_added)) { in AllGatherShardsInternal()
3600 int64* next_channel_id, absl::Span<const int64> selected_dims, in AllReduceAlongShardingDims() argument
3604 b, operand, sharding, next_channel_id, selected_dims, collectives_creator, in AllReduceAlongShardingDims()
[all …]
|