//  [python]
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])
        )
        self.assertTrue(torch.allclose(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]),
        )
        self.assertTrue(torch.allclose(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])
        )
        self.assertTrue(torch.allclose(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)
        self.assertTrue(
            torch.allclose(
                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]),
        )
        self.assertTrue(torch.allclose(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)