code: def test_same_sharding_plan(self) -> None:
      sharding_options = self.enumerator.enumerate(
          module=self.model, sharders=[TWSharder()]
      )
  
      for sharding_option in sharding_options:
          sharding_option.shards[0].perf = Perf(
              fwd_compute=40, fwd_comms=30, bwd_compute=20, bwd_comms=10
          )
          sharding_option.shards[0].storage = Storage(
              hbm=1000 * 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
      }
  
      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
      }
      self.assertEqual(greedy_perf_ranks, memory_balanced_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 TestMemoryBalancedPartitioner obj's underlying member method code with example usage: obj.test_same_sharding_plan() without return types