fix EBTS bug where gen returned less than n total weight (facebook#3316)

Summary:

This could happen if enought `weights` were zero here than the length of `weights` became less than `n`: https://www.internalfb.com/code/fbsource/[b9c286a9f21709eb4c964c3c24bb629b2b218c86]/fbcode/ax/models/discrete/thompson.py?lines=109-113.

The `AlmostEqual` is needed to deal with numerical precision (e.g. 9.0000000002)

Reviewed By: danielcohenlive

Differential Revision: D69253157

Author: sdaulton

Diff for: ax/models/discrete/thompson.py

@@ -104,7 +104,6 @@ def gen(
             objective_weights=objective_weights, outcome_constraints=outcome_constraints
         )
         min_weight = self.min_weight if self.min_weight is not None else 2.0 / k
-
         # Second entry is used for tie-breaking
         weighted_arms = [
             (weights[i], np.random.random(), arms[i])
@@ -128,9 +127,7 @@ def gen(
         if self.uniform_weights:
             top_weights = [1.0 for _ in top_weights]
         else:
-            top_weights = [
-                (x * len(top_weights)) / sum(top_weights) for x in top_weights
-            ]
+            top_weights = [(x * n) / sum(top_weights) for x in top_weights]
         return (
             top_arms,
             top_weights,

Diff for: ax/models/tests/test_eb_thompson.py

@@ -75,7 +75,7 @@ def test_EmpiricalBayesThompsonSamplerGen(self) -> None:
             )
             self.assertEqual(arms, [[4, 4], [3, 3], [2, 2], [1, 1]])
             for weight, expected_weight in zip(
-                weights, [4 * i for i in [0.66, 0.25, 0.07, 0.02]]
+                weights, [5 * i for i in [0.66, 0.25, 0.07, 0.02]]
             ):
                 self.assertAlmostEqual(weight, expected_weight, delta=0.1)
 
@@ -95,7 +95,7 @@ def test_EmpiricalBayesThompsonSamplerWarning(self) -> None:
         )
         self.assertEqual(arms, [[3, 3], [2, 2], [1, 1]])
         for weight, expected_weight in zip(
-            weights, [3 * i for i in [0.74, 0.21, 0.05]]
+            weights, [5 * i for i in [0.74, 0.21, 0.05]]
         ):
             self.assertAlmostEqual(weight, expected_weight, delta=0.1)
 

Diff for: ax/models/tests/test_thompson.py

@@ -105,7 +105,7 @@ def test_ThompsonSamplerMinWeight(self) -> None:
         )
         self.assertEqual(arms, [[4, 4], [3, 3], [2, 2]])
         for weight, expected_weight in zip(
-            weights, [3 * i for i in [0.725, 0.225, 0.05]]
+            weights, [5 * i for i in [0.725, 0.225, 0.05]]
         ):
             self.assertAlmostEqual(weight, expected_weight, 1)