Searched refs:source_dim (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | indexed_array_analysis.h | 186 int64 source_dim() const { return source_dim_; } in source_dim() function
194 explicit ScalarIndexedArray(Array* source, Array* indices, int64 source_dim, in ScalarIndexedArray() argument
198 source_dim_(source_dim), in ScalarIndexedArray()
225 int64 source_dim, in ScalarIndexedConstantArray() argument
228 : ScalarIndexedArray(source, indices, source_dim, in ScalarIndexedConstantArray()
305 ScalarIndexedArray* source, Array* indices, int64 source_dim,
337 Array* source, Array* indices, int64 source_dim, in ConstructScalarIndexedArray() argument
340 return Construct<ScalarIndexedConstantArray>(source, indices, source_dim, in ConstructScalarIndexedArray()
344 return Construct<ScalarIndexedArray>(source, indices, source_dim, in ConstructScalarIndexedArray()
|
| D | indexed_array_analysis.cc | 73 indexed_array->source_dim(), "->[", in ToString()
188 ScalarIndexedArray* source, Array* indices, int64 source_dim, in FoldGatherOfGather() argument
207 EraseAt(&simulated_index, source->source_dim()); in FoldGatherOfGather()
214 EraseAt(&simulated_index, source_dim); in FoldGatherOfGather()
221 FindIndex(source->output_dims(), source_dim); in FoldGatherOfGather()
248 return ConstructScalarIndexedArray(a, inner_indices, source->source_dim(), in FoldGatherOfGather()
290 int64 source_dim = dim_numbers.start_index_map(0); in ComputeArrayForGather() local
299 if (absl::c_linear_search(indexed->output_dims(), source_dim)) { in ComputeArrayForGather()
300 return FoldGatherOfGather(indexed, indices, source_dim, output_dims, in ComputeArrayForGather()
304 return Construct<ScalarIndexedConstantArray>(constant, indices, source_dim, in ComputeArrayForGather()
[all …]
|
| D | sharding_propagation.cc | 756 const int64 source_dim = std::distance(dimensions.begin(), it); in InferShardingFromOperands() local
758 op->sharding().tile_assignment().dim(source_dim)); in InferShardingFromOperands()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 1099 int64 source_dim = source_target_dims[0].first; in ReshardWithAllToAll() local
1101 const int64 group_size = sharding().tile_assignment().dim(source_dim) / in ReshardWithAllToAll()
1111 if (source_dim == j) { in ReshardWithAllToAll()
1133 temp_target_tile_dims[source_dim] = in ReshardWithAllToAll()
1136 sharding().tile_assignment().dim(source_dim); in ReshardWithAllToAll()
1191 (target_dim < source_dim) ? source_dim + 1 : source_dim; in ReshardWithAllToAll()
1209 padded_hlo->shape(), target_dim, source_dim, group_size) in ReshardWithAllToAll()
|