Searched refs:new_source_dim (Results 1 – 2 of 2) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A D | indexed_array_analysis.cc | 519 int64_t new_source_dim = in ReshapeToRemoveDegenerateDims() local
523 new_source, new_indices, new_source_dim, in ReshapeToRemoveDegenerateDims()
574 int64_t new_source_dim = [&]() { in ReshapeToAddDegenerateDims() local
580 if (new_source_shape_dims[new_source_dim] != 1) { in ReshapeToAddDegenerateDims()
589 InsertAt(&new_source_shape_dims, /*index=*/new_source_dim, in ReshapeToAddDegenerateDims()
601 new_source, operand->indices(), new_source_dim, in ReshapeToAddDegenerateDims()
1061 int64_t new_source_dim = dim_numbers.lhs_batch_dimensions_size() + in ComputeArrayForDotWithIndexedLhs() local
1066 new_source, lhs->indices(), new_source_dim, in ComputeArrayForDotWithIndexedLhs()
1097 int64_t new_source_dim = dim_numbers.lhs_batch_dimensions_size() + in ComputeArrayForDotWithIndexedRhs() local
1102 new_source, rhs->indices(), new_source_dim, in ComputeArrayForDotWithIndexedRhs()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| H A D | spmd_partitioner.cc | 1591 int64_t new_source_dim = in ReshardWithAllToAll() local
1598 if (i == new_source_dim) { in ReshardWithAllToAll()
|