code: def test_greedy_partitioner(self) -> None:
      greedy_partitioner = GreedyPerfPartitioner(balance_modules=False)
      balance_modules_greedy_partitioner = GreedyPerfPartitioner(
          balance_modules=True
      )
  
      greedy_sharding_plan = greedy_partitioner.partition(
          proposal=self.sharding_options,
          storage_constraint=self.topology,
      )
      greedy_ranks = {
          sharding_option.name: [
              shard.rank for shard in sharding_option.shards
          ]
          for sharding_option in greedy_sharding_plan
      }
  
      reset_shard_rank(self.sharding_options)
  
      balance_modules_sharding_plan = (
          balance_modules_greedy_partitioner.partition(
              proposal=self.sharding_options,
              storage_constraint=self.topology,
          )
      )
      balance_modules_ranks = {
          sharding_option.name: [
              shard.rank for shard in sharding_option.shards
          ]
          for sharding_option in balance_modules_sharding_plan
      }
  
      greedy_expected_ranks = {
          "weighted_table_0": [0],
          "weighted_table_1": [1],
          "weighted_table_2": [0],
          "table_0": [1],
      }
      balance_modules_expected_ranks = {
          "weighted_table_0": [1],
          "weighted_table_1": [0],
          "weighted_table_2": [1],
          "table_0": [0],
      }
  
      self.assertEqual(greedy_expected_ranks, greedy_ranks)
      self.assertEqual(
          balance_modules_expected_ranks, balance_modules_ranks
      )
  
imports: import copy
  import unittest
  from typing import cast, List
  from unittest.mock import MagicMock
  from torch import nn
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder
  from torchrec.distributed.planner.constants import BATCH_SIZE
  from torchrec.distributed.planner.enumerators import EmbeddingEnumerator
  from torchrec.distributed.planner.partitioners import GreedyPerfPartitioner, MemoryBalancedPartitioner, OrderedDeviceHardware
  from torchrec.distributed.planner.types import DeviceHardware, ParameterConstraints, PartitionByType, Perf, PlannerError, Shard, ShardingOption, Storage, Topology
  from torchrec.distributed.planner.utils import reset_shard_rank
  from torchrec.distributed.test_utils.test_model import TestSparseNN
  from torchrec.distributed.types import ModuleSharder, ShardingType
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
label_desc: Use image node_id 2 for calling the TestBalanceModules obj's underlying member method code with example usage: obj.test_greedy_partitioner() without return types