Merge pull request #176 from JaxGaussianProcesses/remove_chex

Remove chex

Author: thomaspinder

.circleci/config.yml

@@ -1,7 +1,7 @@
 version: 2.1
 
 orbs:
-  python: circleci/[email protected]
+  # python: circleci/[email protected]
   codecov: codecov/[email protected]
 
 commands:
@@ -51,7 +51,6 @@ commands:
       - run:
           name: Upload to PyPI
           command: twine upload dist/* -r << parameters.pkgname >> --verbose
-
   install_pandoc:
     description: "Install pandoc"
     parameters:
@@ -83,15 +82,12 @@ jobs:
     resource_class: large
     steps:
       - checkout
-      - restore_cache:
-          keys:
-            - pip-cache
       - run:
-          name: Update pip
-          command: pip install --upgrade pip
-      - python/install-packages:
-          pkg-manager: pip-dist
-          path-args: .[dev]
+          name: Install dependencies
+          command: |
+            pip install --upgrade pip
+            pip install -r requirements/dev.txt
+            pip install -e .
       - run:
           name: Run tests
           command: |
@@ -103,10 +99,6 @@ jobs:
             curl -Os https://uploader.codecov.io/v0.1.0_4653/linux/codecov
             chmod +x codecov
             ./codecov -t ${CODECOV_TOKEN}
-      - save_cache:
-          key: pip-cache
-          paths:
-            - ~/.cache/pip
       - store_test_results:
           path: test-results
       - store_artifacts:

README.md

@@ -123,7 +123,7 @@ parameter_state = gpx.initialise(posterior, key=key)
 Finally, we run an optimisation loop using the Adam optimiser via the `fit` callable.
 
 ```python
-inference_state = gpx.fit(mll, parameter_state, opt, n_iters=500)
+inference_state = gpx.fit(mll, parameter_state, opt, num_iters=500)
 ```
 
 ## 3. Making predictions

docs/README.md

@@ -41,7 +41,6 @@ description and a code example. The docstring is concluded with a description
 of the objects attributes with corresponding types.
 
 ```python
-@dataclass
 class Prior(AbstractPrior):
     """A Gaussian process prior object. The GP is parameterised by a
     `mean <https://gpjax.readthedocs.io/en/latest/api.html#module-gpjax.mean_functions>`_
@@ -78,9 +77,4 @@ class Prior(AbstractPrior):
 ### Documentation syntax
 
 A helpful cheatsheet for writing restructured text can be found 
-[here](https://github.com/ralsina/rst-cheatsheet/blob/master/rst-cheatsheet.rst). In addition to that, we adopt the following convention when documenting
-`dataclass` objects.
-
-*  Class attributes should be specified using the `Attributes:` tag.
-*  Method argument should be specified using the `Args:` tags.
-*  All attributes and arguments should have types.
+[here](https://github.com/ralsina/rst-cheatsheet/blob/master/rst-cheatsheet.rst).

examples/README.md

@@ -0,0 +1,85 @@
+# Where to find the docs
+
+The GPJax documentation can be found here:
+https://gpjax.readthedocs.io/en/latest/
+
+# How to build the docs
+
+1. Install the requirements using `pip install -r docs/requirements.txt`
+2. Make sure `pandoc` is installed
+3. Run the make script `make html`
+
+The corresponding HTML files can then be found in `docs/_build/html/`.
+
+# How to write code documentation
+
+Our documentation it is written in ReStructuredText for Sphinx. This is a
+meta-language that is compiled into online documentation. For more details see
+[Sphinx's documentation](https://www.sphinx-doc.org/en/master/usage/restructuredtext/index.html).
+As a result, our docstrings adhere to a specific syntax that has to be kept in
+mind. Below we provide some guidelines.
+
+## How much information to put in a docstring
+
+A docstring should be informative. If in doubt, then it is best to add more
+information to a docstring than less. Many users will skim documentation, so
+please ensure the opening sentence or two of a docstring contains the core
+information. Adding examples and mathematical descriptions to documentation is
+highly desirable.
+
+We are making an active effort within GPJax to improve our documentation. If you
+spot any areas where there is missing information within the existing
+documentation, then please either raise an issue or 
+[create a pull request](https://gpjax.readthedocs.io/en/latest/contributing.html).
+
+## An example docstring
+
+An example docstring that adheres the principles of GPJax is given below. 
+The docstring contains a simple, snappy introduction with links to auxillary 
+components. More detail is then provided in the form of a mathematical 
+description and a code example. The docstring is concluded with a description
+of the objects attributes with corresponding types.
+
+```python
+class Prior(AbstractPrior):
+    """A Gaussian process prior object. The GP is parameterised by a
+    `mean <https://gpjax.readthedocs.io/en/latest/api.html#module-gpjax.mean_functions>`_
+    and `kernel <https://gpjax.readthedocs.io/en/latest/api.html#module-gpjax.kernels>`_ function.
+
+    A Gaussian process prior parameterised by a mean function :math:`m(\\cdot)` and a kernel
+    function :math:`k(\\cdot, \\cdot)` is given by
+
+    .. math::
+
+        p(f(\\cdot)) = \mathcal{GP}(m(\\cdot), k(\\cdot, \\cdot)).
+
+    To invoke a ``Prior`` distribution, only a kernel function is required. By default,
+    the mean function will be set to zero. In general, this assumption will be reasonable
+    assuming the data being modelled has been centred.
+
+    Example:
+        >>> import gpjax as gpx
+        >>>
+        >>> kernel = gpx.kernels.RBF()
+        >>> prior = gpx.Prior(kernel = kernel)
+
+    Attributes:
+        kernel (Kernel): The kernel function used to parameterise the prior.
+        mean_function (MeanFunction): The mean function used to parameterise the prior. Defaults to zero.
+        name (str): The name of the GP prior. Defaults to "GP prior".
+    """
+
+    kernel: Kernel
+    mean_function: Optional[AbstractMeanFunction] = Zero()
+    name: Optional[str] = "GP prior"
+```
+
+### Documentation syntax
+
+A helpful cheatsheet for writing restructured text can be found 
+[here](https://github.com/ralsina/rst-cheatsheet/blob/master/rst-cheatsheet.rst). In addition to that, we adopt the following convention when documenting
+`` objects.
+
+*  Class attributes should be specified using the `Attributes:` tag.
+*  Method argument should be specified using the `Args:` tags.
+*  All attributes and arguments should have types.

examples/barycentres.pct.py

@@ -115,7 +115,7 @@ def fit_gp(x: jax.Array, y: jax.Array) -> dx.MultivariateNormalTri:
         objective=negative_mll,
         parameter_state=parameter_state,
         optax_optim=optimiser,
-        n_iters=1000,
+        num_iters=1000,
     )
 
     learned_params, training_history = inference_state.unpack()

examples/classification.pct.py

@@ -91,7 +91,7 @@
     objective=negative_mll,
     parameter_state=parameter_state,
     optax_optim=optimiser,
-    n_iters=1000,
+    num_iters=1000,
 )
 
 map_estimate, training_history = inference_state.unpack()

examples/collapsed_vi.pct.py

@@ -109,7 +109,7 @@
     objective=negative_elbo,
     parameter_state=parameter_state,
     optax_optim=optimiser,
-    n_iters=2000,
+    num_iters=2000,
 )
 
 learned_params, training_history = inference_state.unpack()

examples/graph_kernels.pct.py

@@ -137,7 +137,7 @@
     objective=negative_mll,
     parameter_state=parameter_state,
     optax_optim=optimiser,
-    n_iters=1000,
+    num_iters=1000,
 )
 
 learned_params, training_history = inference_state.unpack()

examples/haiku.pct.py

@@ -185,7 +185,7 @@ def forward(x):
     objective=negative_mll,
     parameter_state=parameter_state,
     optax_optim=optimiser,
-    n_iters=2500,
+    num_iters=2500,
 )
 
 learned_params, training_history = inference_state.unpack()

examples/kernels.pct.py

@@ -228,28 +228,7 @@ def _initialise_params(self, key: jr.PRNGKey) -> dict:
 # domain is a circle, this is $2\pi$. Next we define the kernel's `__call__`
 # function which is a direct implementation of Equation (1). Finally, we define
 # the Kernel's parameter property which contains just one value $\tau$ that we
-# initialise to 4 in the kernel's `__post_init__`.
-#
-# #### Aside on dataclasses
-#
-# One can see in the above definition of a `Polar` kernel that we decorated the
-# class with a `@dataclass` command. Dataclasses are simply regular classs
-# objects in Python, however, much of the boilerplate code has been removed. For
-# example, without a `@dataclass` decorator, the instantiation of the above
-# `Polar` kernel would be done through
-# ```python
-# class Polar(jk.kernels.AbstractKernel):
-#     def __init__(self, period: float = 2*jnp.pi):
-#         super().__init__()
-#         self.period = period
-# ```
-# As objects become increasingly large and complex, the conciseness of a
-# dataclass becomes increasingly attractive. To ensure full compatability with
-# Jax, it is crucial that the dataclass decorator is imported from Chex, not
-# base Python's `dataclass` module. Functionally, the two objects are identical.
-# However, unlike regular Python dataclasses, it is possilbe to apply operations
-# such as `jit`, `vmap` and `grad` to the dataclasses given by Chex as they are
-# registrered PyTrees.
+# initialise to 4 in the kernel's `__init__`.
 #
 #
 # ### Custom Parameter Bijection
@@ -312,7 +291,7 @@ def _initialise_params(self, key: jr.PRNGKey) -> dict:
     objective=negative_mll,
     parameter_state=parameter_state,
     optax_optim=optimiser,
-    n_iters=1000,
+    num_iters=1000,
 )
 
 learned_params, training_history = inference_state.unpack()