code: def _group_and_sort_non_uniform_sharding_options(
      sharding_options: List[ShardingOption],
      sort_by: SortBy = SortBy.STORAGE,
      balance_modules: bool = False,
  ) -> List[ShardingOptionGroup]:
      # count modules by name
      param_count: Dict[str, int] = {}
      for sharding_option in sharding_options:
          path = sharding_option.path
          if path not in param_count:
              param_count[path] = 0
          param_count[path] += 1
      logger.debug(f"param_count is {param_count}")
  
      sharding_option_groups_by_dependency = {}
      for sharding_option in sharding_options:
          if (
              sharding_option.partition_by
              == PartitionByType.UNIFORM.value
          ):
              continue
  
          group_key = sharding_option.dependency or sharding_option.fqn
          if group_key not in sharding_option_groups_by_dependency:
              sharding_option_groups_by_dependency[
                  group_key
              ] = ShardingOptionGroup(
                  [sharding_option],
                  sharding_option.total_storage,
                  sharding_option.total_perf,
                  # negative value to indicate that smaller modules should be sorted first
                  param_count=-param_count[sharding_option.path],
              )
          else:
              sharding_option_groups_by_dependency[
                  group_key
              ].sharding_options.append(sharding_option)
              sharding_option_groups_by_dependency[
                  group_key
              ].storage_sum += sharding_option.total_storage
              sharding_option_groups_by_dependency[
                  group_key
              ].perf_sum += sharding_option.total_perf
  
      sharding_option_groups = list(
          sharding_option_groups_by_dependency.values()
      )
  
      sort_by_attributes: List[str] = []
      if balance_modules:
          sort_by_attributes.append("param_count")
  
      if sort_by == SortBy.STORAGE:
          sort_by_attributes.append("storage_sum")
      elif sort_by == SortBy.PERF:
          sort_by_attributes.append("perf_sum")
      else:
          raise RuntimeError(f"Unexpected sort_by: {sort_by}")
  
      sharding_option_groups.sort(
          key=lambda group: [
              getattr(group, attr) for attr in sort_by_attributes
          ],
          reverse=True,
      )
  
      return sharding_option_groups
  
imports: import copy
  import heapq
  import logging
  from dataclasses import dataclass
  from enum import Enum
  from typing import cast, Dict, List, Optional
  from torchrec.distributed.planner.perf_models import NoopPerfModel
  from torchrec.distributed.planner.types import DeviceHardware, PartitionByType, Partitioner, Perf, PerfModel, PlannerError, PlannerErrorType, ShardingOption, Storage, Topology
  from torchrec.distributed.planner.utils import bytes_to_gb, reset_shard_rank
  from torchrec.distributed.types import ShardingType
label_desc: Use image node_id 3 for calling a global function with example usage: _group_and_sort_non_uniform_sharding_options(sharding_options, sort_by, balance_modules) and returns: sharding_option_groups