Searched defs:ag_shape (Results 1 – 2 of 2) sorted by relevance
1659 int64 channel_id, int64 all_gather_dimension) { in GetPerGroupCollectiveOpsCreator()
3589 auto ag_shape = operand->shape(); in AllGatherShardsInternal() local