Searched refs:new_tile (Results 1 – 2 of 2) sorted by relevance
112 Array<int64> new_tile(new_tile_dims); in MergeSharding() local135 new_tile.Each([&](absl::Span<const int64> indices, int64* device) { in MergeSharding()187 new_tile.Reshape(new_tile_dims); in MergeSharding()188 *to_merge = HloSharding::Tile(new_tile, merged_metadata); in MergeSharding()190 *to_merge = HloSharding::PartialTile(new_tile, merged_metadata); in MergeSharding()1168 auto new_tile = transposed.tile_assignment(); in PartiallyReplicateTiledShardingOnDims() local1180 new_tile.Reshape(new_tile_shape); in PartiallyReplicateTiledShardingOnDims()1181 return HloSharding::PartialTile(new_tile, sharding.metadata()); in PartiallyReplicateTiledShardingOnDims()1199 auto new_tile = sharding.tile_assignment(); in RemoveShapeDimensions() local1200 new_tile.Reshape(new_tile_shape); in RemoveShapeDimensions()[all …]
2115 auto new_tile = rhs.sharding().tile_assignment(); in PartitionDotGroupOnContracting() local2116 new_tile.Reshape(rhs_tile_shape); in PartitionDotGroupOnContracting()2118 ? HloSharding::PartialTile(new_tile) in PartitionDotGroupOnContracting()2119 : HloSharding::Tile(new_tile); in PartitionDotGroupOnContracting()2124 auto new_tile = lhs.sharding().tile_assignment(); in PartitionDotGroupOnContracting() local2125 new_tile.Reshape(lhs_tile_shape); in PartitionDotGroupOnContracting()2127 ? HloSharding::PartialTile(new_tile) in PartitionDotGroupOnContracting()2128 : HloSharding::Tile(new_tile); in PartitionDotGroupOnContracting()