Create group assignment for XLA cross replica ops (physical pnums).

def _create_group_assignment(self, mesh_axes):
    """Create group assignment for XLA cross replica ops (physical pnums)."""

    partitioning = {}
    for logical_pnum in xrange(self.size):
      group = mtf.pnum_to_group(self.shape, mesh_axes, logical_pnum)
      if group not in partitioning:
        partitioning[group] = []
      partitioning[group].append(self.l2p(logical_pnum))
    group_assignment = []
    for group, physical_pnums in partitioning.items():
      group_assignment.append(physical_pnums)
    return group_assignment