code: def _create_group_assignment(self, mesh_axes):
      
  
      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
docs: Create group assignment for XLA cross replica ops (physical pnums).