code: def test_different_sharding_plan(self) -> None:
      sharding_options = self.enumerator.enumerate(
          module=self.model, sharders=[TWSharder()]
      )
  
      for i, sharding_option in enumerate(sharding_options):
          sharding_option.shards[0].perf = Perf(
              fwd_compute=40 * (i + 1),
              fwd_comms=0,
              bwd_compute=0,
              bwd_comms=0,
          )
          sharding_option.shards[0].storage = Storage(
              hbm=(1500 - i * 500) * 1024**2, ddr=1000 * 1024**2
          )
  
      greedy_perf_sharding_plan = (
          self.greedy_perf_partitioner.partition(
              proposal=sharding_options,
              storage_constraint=self.topology,
          )
      )
      greedy_perf_ranks = {
          sharding_option.name: [
              shard.rank for shard in sharding_option.shards
          ]
          for sharding_option in greedy_perf_sharding_plan
      }
      greedy_perf_expected_ranks = {
          "table_0": [0],
          "table_1": [1],
          "table_2": [0],
      }
      self.assertEqual(greedy_perf_ranks, greedy_perf_expected_ranks)
  
      greedy_perf_hbm_uses = [0] * self.topology.world_size
      for sharding_option in sharding_options:
          for shard in sharding_option.shards:
              if shard.storage and shard.rank is not None:
                  greedy_perf_hbm_uses[
                      shard.rank
                  ] += shard.storage.hbm  # pyre-ignore[16]
  
      reset_shard_rank(sharding_options)
      memory_balanced_sharding_plan = (
          self.memory_balanced_partitioner.partition(
              proposal=sharding_options,
              storage_constraint=self.topology,
          )
      )
      memory_balanced_ranks = {
          sharding_option.name: [
              shard.rank for shard in sharding_option.shards
          ]
          for sharding_option in memory_balanced_sharding_plan
      }
      memory_balanced_expected_ranks = {
          "table_0": [0],
          "table_1": [1],
          "table_2": [0],
      }
      self.assertEqual(
          memory_balanced_ranks, memory_balanced_expected_ranks
      )
  
      memory_balanced_hbm_uses = [0.0] * self.topology.world_size
      for sharding_option in sharding_options:
          for shard in sharding_option.shards:
              if shard.storage and shard.rank:
                  memory_balanced_hbm_uses[
                      shard.rank
                  ] += shard.storage.hbm
  
      self.assertTrue(
          max(memory_balanced_hbm_uses) < max(greedy_perf_hbm_uses)
      )
  
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 3 for calling the TestMemoryBalancedPartitioner obj's underlying member method code with example usage: obj.test_different_sharding_plan() without return types