class_docstr: Memory balanced Partitioner.
  
  Args:
      max_search_count (int): Maximum number of times to call the
          GreedyPartitioner.
      tolerance (float): The maximum acceptable difference between the
          original plan and the new plan. If tolerance is 1, that means a new
          plan will be rejected if its perf is 200% of the original plan
          (i.e., the plan is 100% worse).
      balance_modules (bool): Whether to sort by modules first, where
          smaller modules will be sorted first. In effect, this will place
          tables in each module in a balanced way.
code: def __init__(
      self,
      max_search_count: int = 10,
      tolerance: float = 0.02,
      balance_modules: bool = False,
  ) -> None:
      self._max_search_count: int = max_search_count
      self._tolerance: float = tolerance
      self._balance_modules: bool = balance_modules
  
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 1 to create a new MemoryBalancedPartitioner object from inherited base classes: Partitioner with example: obj = MemoryBalancedPartitioner(max_search_count, tolerance, balance_modules)