Searched refs:update_or_gather_sharding (Results 1 – 1 of 1) sorted by relevance
846 const Shape& operand_shape, const HloSharding& update_or_gather_sharding, in PassthroughGatherOutputOrScatterUpdateToOperand() argument851 if (update_or_gather_sharding.IsTileMaximal()) { in PassthroughGatherOutputOrScatterUpdateToOperand()852 return update_or_gather_sharding; in PassthroughGatherOutputOrScatterUpdateToOperand()864 update_or_gather_sharding.tile_assignment().dim(offset_dim); in PassthroughGatherOutputOrScatterUpdateToOperand()876 if (update_or_gather_sharding.ReplicateOnLastTileDim()) { in PassthroughGatherOutputOrScatterUpdateToOperand()878 update_or_gather_sharding.tile_assignment().dimensions().back()); in PassthroughGatherOutputOrScatterUpdateToOperand()880 Array<int64> tile_assignment = update_or_gather_sharding.tile_assignment(); in PassthroughGatherOutputOrScatterUpdateToOperand()885 return update_or_gather_sharding.ReplicateOnLastTileDim() in PassthroughGatherOutputOrScatterUpdateToOperand()887 update_or_gather_sharding.metadata()) in PassthroughGatherOutputOrScatterUpdateToOperand()889 update_or_gather_sharding.metadata()); in PassthroughGatherOutputOrScatterUpdateToOperand()