Searched defs:new_tile_assignment (Results 1 – 4 of 4) sorted by relevance
420 Array<int64> new_tile_assignment = sharding.tile_assignment(); in ReshapeSharding() local438 Array<int64> new_tile_assignment(sharding.tile_assignment().dimensions()); in ReverseSharding() local544 Array<int64> new_tile_assignment = index_sharding.tile_assignment(); in GatherOutputSharding() local595 Array<int64> new_tile_assignment = relevant_output_sharding.tile_assignment(); in GatherIndexSharding() local689 Array<int64> new_tile_assignment = data_sharding.tile_assignment(); in ScatterIndexSharding() local722 Array<int64> new_tile_assignment = index_sharding.tile_assignment(); in ScatterDataSharding() local
782 Array<int64> new_tile_assignment = op->sharding().tile_assignment(); in InferShardingFromOperands() local1211 auto new_tile_assignment = in GetShardingFromUser() local
61 auto new_tile_assignment = Array<int64>(new_tile_dims); in PartialTile() local
1954 auto new_tile_assignment = to_adjust->tile_assignment(); in GatherOperandsShardedAcrossParallelDims() local