Cython Sobol --> PyTorch Sobol! (facebook#165)

Summary:
Pull Request resolved: facebook#165

Goodbye Cython! Goodbye multiple wheels!

queen_of_hearts

Differential Revision: D17241392

fbshipit-source-id: 69ea410604af93ba16381babd093a97bbb5d9297

.travis.yml

@@ -12,7 +12,7 @@ jobs:
       stage: Test
       python: "3.6"
       install:
-        - pip install cython numpy
+        - pip install numpy
         - pip install git+https://github.com/cornellius-gp/gpytorch.git
         - pip install git+https://github.com/pytorch/botorch.git
         - pip install -q -e .[dev,mysql,notebook]
@@ -24,7 +24,7 @@ jobs:
       python: "3.7"
       dist: xenial
       install:
-        - pip install cython numpy
+        - pip install numpy
         - pip install git+https://github.com/cornellius-gp/gpytorch.git
         - pip install git+https://github.com/pytorch/botorch.git
         - pip install -q -e .[dev,mysql,notebook]
@@ -46,7 +46,7 @@ jobs:
     - name: "Docs: Sphinx (Python 3.6)"
       python: "3.6"
       install:
-        - pip install cython numpy
+        - pip install numpy
         - pip install git+https://github.com/cornellius-gp/gpytorch.git
         - pip install git+https://github.com/pytorch/botorch.git
         - pip install -q -e .[dev,mysql,notebook,gpy]
@@ -59,7 +59,7 @@ jobs:
       python: "3.7"
       env: ALLOW_FAILURE=true
       install:
-        - pip install cython numpy
+        - pip install numpy
         - pip install botorch
         - pip install -q -e .[dev,mysql,notebook]
       script:
@@ -73,7 +73,7 @@ jobs:
     - stage: Deploy Website
       python: "3.6"
       install:
-        - pip install cython numpy
+        - pip install numpy
         - pip install git+https://github.com/cornellius-gp/gpytorch.git
         - pip install git+https://github.com/pytorch/botorch.git
         - pip install -q -e .[dev,mysql,notebook]

README.md

@@ -111,7 +111,6 @@ pip3 install SQLAlchemy
 You can install the latest (bleeding edge) version from Git:
 
 ```
-pip3 install cython numpy  # needed for compiling Cython code
 pip3 install git+ssh://git@github.com/facebook/Ax.git#egg=Ax
 ```
 
@@ -146,7 +145,7 @@ When contributing to Ax, we recommend cloning the [repository](https://github.co
 pip3 install git+https://github.com/cornellius-gp/gpytorch.git
 pip3 install git+https://github.com/pytorch/botorch.git
 
-pip3 install cython numpy  # needed for compiling Cython code
+pip3 install numpy  # needed for compiling Cython code
 git clone https://github.com/facebook/ax.git
 cd ax
 pip3 install -e .[notebook,mysql,dev]

ax/models/random/sobol.py

@@ -9,7 +9,8 @@
 from ax.models.model_utils import tunable_feature_indices
 from ax.models.random.base import RandomModel
 from ax.utils.common.docutils import copy_doc
-from ax.utils.stats.sobol import SobolEngine  # pyre-ignore: Not handling .pyx properly
+from ax.utils.common.typeutils import not_none
+from torch.quasirandom import SobolEngine
 
 
 class SobolGenerator(RandomModel):
@@ -29,7 +30,6 @@ class SobolGenerator(RandomModel):
 
     """
 
-    # pyre-fixme[31]: Expression `SobolEngine` is not a valid type.
     engine: Optional[SobolEngine] = None
 
     def __init__(
@@ -45,7 +45,6 @@ def __init__(
         # Initialize engine on gen.
         self._engine = None
 
-    # pyre-fixme[11]: Type `SobolEngine` is not defined.
     def init_engine(self, n_tunable_features: int) -> SobolEngine:
         """Initialize singleton SobolEngine, only on gen.
 
@@ -58,13 +57,13 @@ def init_engine(self, n_tunable_features: int) -> SobolEngine:
 
         """
         if not self._engine:
-            self._engine = SobolEngine(  # pyre-ignore: .pyx not parsed properly.
-                dimen=n_tunable_features, scramble=self.scramble, seed=self.seed
+            self._engine = SobolEngine(
+                dimension=n_tunable_features, scramble=self.scramble, seed=self.seed
             ).fast_forward(self.init_position)
         return self._engine
 
     @property
-    def engine(self) -> Optional[SobolEngine]:  # pyre-fixme[31]: not valid type.
+    def engine(self) -> Optional[SobolEngine]:
         """Return a singleton SobolEngine."""
         return self._engine
 
@@ -111,7 +110,7 @@ def gen(
             rounding_func=rounding_func,
         )
         if self.engine:
-            self.init_position = self.engine.num_generated
+            self.init_position = not_none(self.engine).num_generated
         return (points, weights)
 
     @copy_doc(Model._get_state)
@@ -129,7 +128,7 @@ def _gen_samples(self, n: int, tunable_d: int) -> np.ndarray:
                 particular value.
         """
         if self.engine is None:
-            raise ValueError(  # pragma nocover
+            raise ValueError(  # pragma: no cover
                 "Sobol Engine must be initialized before candidate generation."
             )
-        return self.engine.draw(n)
+        return not_none(self.engine).draw(n)