code: def allocate_budget(
      model: torch.Tensor,
      clfs: torch.Tensor,
      budget: int,
      allocation_priority: torch.Tensor,
  ) -> torch.Tensor:
      # min size is size of table at 0 CLF
      min_size_bytes = EmbeddingOffloadScaleupProposer.clf_to_bytes(
          model, 0
      )
      max_size_bytes = EmbeddingOffloadScaleupProposer.clf_to_bytes(
          model, 1
      )
      table_size_bytes = EmbeddingOffloadScaleupProposer.clf_to_bytes(
          model, clfs
      )
      cache_size_bytes = table_size_bytes - min_size_bytes
      max_cache_size_bytes = max_size_bytes - min_size_bytes
  
      # We have budget bytes to share across the tables in. We want to increase the
      # cache_size_bytes of each table in proportion to their allocation priority
      # fraction. If we raise the cache_size_bytes to beyond max_cache_size_bytes,
      # this is equivalent to reaching CLF=1.0, so we clip the memory to 1.0, and
      # reassign the released budget in a subsequent pass.
      num_pass = 0
      while budget > 1 and num_pass < 128:
          num_pass += 1
          # mask is False for tables at >= max_size, and True otherwise. This allows
          # us to remove tables that have already reached full size in one round from
          # being dealt more budget in subsequent rounds.
          mask = (min_size_bytes + cache_size_bytes) < max_size_bytes
          if mask.sum() == 0:
              break
  
          logging.debug(
              f"[allocate_budget] pass={num_pass}, budget={budget}, #cache_tables={mask.sum()}"
          )
  
          # switch to 64bit float to avoid rounding errors, as table cache sizes can
          # easily be > 2^24.
          masked_priority = (mask * allocation_priority).to(
              torch.float64
          )
          increase_ratio = masked_priority / torch.sum(masked_priority)
          proposed_increase_bytes = budget * increase_ratio
          new_cache_size_bytes = torch.minimum(
              cache_size_bytes + proposed_increase_bytes,
              max_cache_size_bytes,
          )
          actual_increase_bytes = (
              new_cache_size_bytes - cache_size_bytes
          )
  
          budget -= torch.sum(actual_increase_bytes)
          cache_size_bytes = new_cache_size_bytes
          # TODO: consider trade off of using remaining budget to push >0.95 tables
          # to HBM vs spending that budget on improving hit rate on other tables in
          # next pass.
  
      # cache_size_bytes are the new cache sizes we want to use. We convert them back
      # to clfs by dividing by max_cache_size_bytes, which has isolated the clf
      # portion of the table size from the fixed overheads.
      # convert 64bit values back to original clf precision
      return (cache_size_bytes / max_cache_size_bytes).to(clfs.dtype)
  
imports: import copy
  import itertools
  import logging
  from decimal import Decimal
  from typing import cast, Dict, List, Optional, Set, Tuple, Union
  import torch
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel
  from torchrec.distributed.planner.types import Enumerator, Perf, Proposer, ShardingOption, Topology
  from torchrec.distributed.planner.utils import BinarySearchPredicate, bytes_to_gb, prod
label_desc: Use image node_id 12 for calling the EmbeddingOffloadScaleupProposer obj's underlying member method code with example usage: obj.allocate_budget(model, clfs, budget, allocation_priority) and returns: unknown