Home
last modified time | relevance | path

Searched refs:reshape_tile_assignment (Results 1 – 2 of 2) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc368 auto reshape_tile_assignment = partial_sharding.tile_assignment(); in PartialReplicateReshardCompatibleSharding() local
369 reshape_tile_assignment.Reshape(reshape_dimensions); in PartialReplicateReshardCompatibleSharding()
382 ? HloSharding::PartialTile(reshape_tile_assignment) in PartialReplicateReshardCompatibleSharding()
383 : HloSharding::Tile(reshape_tile_assignment), in PartialReplicateReshardCompatibleSharding()
Dspmd_partitioner.cc1279 auto reshape_tile_assignment = partial_replicate_sharding.tile_assignment(); in ReshardPartialReplicateWithAllToAll() local
1280 reshape_tile_assignment.Reshape(tile_sharding.tile_assignment().dimensions()); in ReshardPartialReplicateWithAllToAll()
1281 if (reshape_tile_assignment != tile_sharding.tile_assignment()) { in ReshardPartialReplicateWithAllToAll()