code: def test_allocate_budget(self) -> None:
      model = torch.tensor([[1.0, 0.0], [2.0, 3.0], [4.0, 5.0]])
      got = EmbeddingOffloadScaleupProposer.clf_to_bytes(
          model, torch.tensor([0, 0.5, 1])
      )
      torch.testing.assert_close(got, torch.tensor([0, 4, 9]))
  
      # Scenario 1, enough budget to scale everything to 1.0
      model = torch.tensor(
          [
              [30_000_000, 2_000_000],
              [30_000_000, 2_000_000],
              [30_000_000, 2_000_000],
          ]
      )
      mins = torch.tensor([0.1, 0.1, 1])
      budget = 100_000_000
      got = EmbeddingOffloadScaleupProposer.allocate_budget(
          model,
          clfs=torch.tensor(mins),
          budget=budget,
          allocation_priority=torch.tensor([2, 2, 2]),
      )
      torch.testing.assert_close(got, torch.tensor([1.0, 1.0, 1.0]))
      increase = (
          EmbeddingOffloadScaleupProposer.clf_to_bytes(model, got).sum()
          - EmbeddingOffloadScaleupProposer.clf_to_bytes(
              model, mins
          ).sum()
      ).item()
      self.assertLess(increase, budget)
  
      # Scenario 2, limited budget, uniform scale up
      model = torch.tensor(
          [
              [30_000_000, 2_000_000],
              [30_000_000, 2_000_000],
              [30_000_000, 2_000_000],
          ]
      )
      mins = torch.tensor([0.1, 0.1, 1])
      budget = 10_000_000
      got = EmbeddingOffloadScaleupProposer.allocate_budget(
          model,
          clfs=mins,
          budget=budget,
          allocation_priority=torch.tensor([2, 2, 2]),
      )
      torch.testing.assert_close(
          got, torch.tensor([0.26667, 0.26667, 1.0])
      )
      increase = (
          EmbeddingOffloadScaleupProposer.clf_to_bytes(model, got).sum()
          - EmbeddingOffloadScaleupProposer.clf_to_bytes(
              model, mins
          ).sum()
      )
      self.assertEqual(increase, budget)
  
      # Scenario 3, limited budget, skewed scale up
      model = torch.tensor(
          [
              [30_000_000, 2_000_000],
              [30_000_000, 2_000_000],
              [30_000_000, 2_000_000],
          ]
      )
      mins = torch.tensor([0.1, 0.1, 1])
      budget = 10_000_000
      got = EmbeddingOffloadScaleupProposer.allocate_budget(
          model,
          clfs=mins,
          budget=budget,
          allocation_priority=torch.tensor([2, 4, 2]),
      )
      # increase is twice as much for table 2 (started at 0.1)
      torch.testing.assert_close(
          got, torch.tensor([0.1 + 0.11111, 0.1 + 2 * 0.11111, 1.0])
      )
      increase = (
          EmbeddingOffloadScaleupProposer.clf_to_bytes(model, got).sum()
          - EmbeddingOffloadScaleupProposer.clf_to_bytes(
              model, mins
          ).sum()
      )
      self.assertEqual(increase, budget)
  
      # Scenario 4, multi-pass scale up
      model = torch.tensor(
          [
              [30_000_000, 2_000_000],
              [30_000_000, 2_000_000],
              [30_000_000, 2_000_000],
          ]
      )
      mins = torch.tensor([0.1, 0.3, 0.5])
      budget = 50_000_000
      got = EmbeddingOffloadScaleupProposer.allocate_budget(
          model,
          clfs=mins,
          budget=budget,
          allocation_priority=torch.tensor([1, 2, 100]),
      )
      torch.testing.assert_close(got, torch.tensor([0.56667, 1.0, 1.0]))
      increase = (
          EmbeddingOffloadScaleupProposer.clf_to_bytes(model, got).sum()
          - EmbeddingOffloadScaleupProposer.clf_to_bytes(
              model, mins
          ).sum()
      )
      self.assertEqual(increase, budget)
  
imports: import unittest
  from typing import cast, List, Optional
  from unittest.mock import MagicMock
  import torch
  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.proposers import EmbeddingOffloadScaleupProposer, GreedyProposer, GridSearchProposer, proposers_to_proposals_list, UniformProposer
  from torchrec.distributed.planner.types import Enumerator, ParameterConstraints, Proposer, ShardingOption, Topology
  from torchrec.distributed.test_utils.test_model import TestSparseNN
  from torchrec.distributed.types import CacheParams, CacheStatistics, ModuleSharder, ShardingType
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
label_desc: Use image node_id 5 for calling the TestProposers obj's underlying member method code with example usage: obj.test_allocate_budget() without return types