From 13cbd31bf3b7fbd6f7bf596628fcee42f94476d8 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Mon, 11 Oct 2021 13:26:35 +0100
Subject: [PATCH 01/23] added quadrature Fourier features and restructured
 module

---
 gpflux/layers/basis_functions/__init__.py     |   5 +-
 .../fourier_features/__init__.py              |  25 ++++
 .../fourier_features/quadrature.py            | 107 ++++++++++++++++++
 .../random.py}                                |  78 ++++++-------
 .../{ => fourier_features}/utils.py           |  30 +++--
 .../test_random_fourier_features.py           |   4 +-
 tests/gpflux/sampling/test_sample.py          |   2 +-
 7 files changed, 189 insertions(+), 62 deletions(-)
 create mode 100644 gpflux/layers/basis_functions/fourier_features/__init__.py
 create mode 100644 gpflux/layers/basis_functions/fourier_features/quadrature.py
 rename gpflux/layers/basis_functions/{random_fourier_features.py => fourier_features/random.py} (83%)
 rename gpflux/layers/basis_functions/{ => fourier_features}/utils.py (83%)

diff --git a/gpflux/layers/basis_functions/__init__.py b/gpflux/layers/basis_functions/__init__.py
index ce7890de..57bdf35f 100644
--- a/gpflux/layers/basis_functions/__init__.py
+++ b/gpflux/layers/basis_functions/__init__.py
@@ -14,8 +14,5 @@
 # limitations under the License.
 #
 """
-A kernel's features for efficient sampling, used by
-:class:`gpflux.sampling.KernelWithFeatureDecomposition`
+Basis functions.
 """
-
-from gpflux.layers.basis_functions.random_fourier_features import RandomFourierFeatures
diff --git a/gpflux/layers/basis_functions/fourier_features/__init__.py b/gpflux/layers/basis_functions/fourier_features/__init__.py
new file mode 100644
index 00000000..989028b1
--- /dev/null
+++ b/gpflux/layers/basis_functions/fourier_features/__init__.py
@@ -0,0 +1,25 @@
+#
+# Copyright (c) 2021 The GPflux Contributors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""
+A kernel's features for efficient sampling, used by
+:class:`gpflux.sampling.KernelWithFeatureDecomposition`
+"""
+
+from gpflux.layers.basis_functions.fourier_features.quadrature import QuadratureFourierFeatures
+from gpflux.layers.basis_functions.fourier_features.random import (
+    RandomFourierFeatures,
+    RandomFourierFeaturesCosine,
+)
diff --git a/gpflux/layers/basis_functions/fourier_features/quadrature.py b/gpflux/layers/basis_functions/fourier_features/quadrature.py
new file mode 100644
index 00000000..6a4f4782
--- /dev/null
+++ b/gpflux/layers/basis_functions/fourier_features/quadrature.py
@@ -0,0 +1,107 @@
+#
+# Copyright (c) 2021 The GPflux Contributors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+""" A kernel's features and coefficients using quadrature Fourier features (QFF). """
+
+from typing import Mapping
+
+import tensorflow as tf
+
+import gpflow
+from gpflow.base import TensorType
+from gpflow.quadrature.gauss_hermite import ndgh_points_and_weights
+
+from gpflux.layers.basis_functions.fourier_features.utils import (
+    QFF_SUPPORTED_KERNELS,
+    _mapping_concat,
+)
+from gpflux.types import ShapeType
+
+
+class QuadratureFourierFeatures(tf.keras.layers.Layer):
+    def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
+        """
+        :param kernel: kernel to approximate using a set of random features.
+        :param output_dim: total number of basis functions used to approximate
+            the kernel.
+        """
+        super().__init__(**kwargs)
+
+        assert isinstance(kernel, QFF_SUPPORTED_KERNELS), "Unsupported Kernel"
+        self.kernel = kernel
+        self.n_components = n_components  # M: number of quadrature points
+        if kwargs.get("input_dim", None):
+            self._input_dim = kwargs["input_dim"]
+            self.build(tf.TensorShape([self._input_dim]))
+        else:
+            self._input_dim = None
+
+    def build(self, input_shape: ShapeType) -> None:
+        """
+        Creates the variables of the layer.
+        See `tf.keras.layers.Layer.build()
+        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#build>`_.
+        """
+        input_dim = input_shape[-1]
+        abscissa_value, weights_value = ndgh_points_and_weights(
+            dim=input_dim, n_gh=self.n_components
+        )
+
+        # Quadrature node points
+        self.abscissa = tf.Variable(initial_value=abscissa_value, trainable=False)
+        # Gauss-Hermite weights (not to be confused with random Fourier feature
+        # weights or NN weights)
+        self.weights = tf.Variable(initiial_value=weights_value, trainable=False)
+        super().build(input_shape)
+
+    def compute_output_shape(self, input_shape: ShapeType) -> tf.TensorShape:
+        """
+        Computes the output shape of the layer.
+        See `tf.keras.layers.Layer.compute_output_shape()
+        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#compute_output_shape>`_.
+        """
+        # TODO: Keras docs say "If the layer has not been built, this method
+        # will call `build` on the layer." -- do we need to do so?
+        input_dim = input_shape[-1]
+        output_dim = 2 * self.n_components ** input_dim
+        tensor_shape = tf.TensorShape(input_shape).with_rank(2)
+        return tensor_shape[:-1].concatenate(output_dim)
+
+    def get_config(self) -> Mapping:
+        """
+        Returns the config of the layer.
+        See `tf.keras.layers.Layer.get_config()
+        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#get_config>`_.
+        """
+        config = super().get_config()
+        config.update(
+            {"kernel": self.kernel, "n_components": self.n_components, "input_dim": self._input_dim}
+        )
+
+        return config
+
+    def call(self, inputs: TensorType) -> tf.Tensor:
+        """
+        Evaluate the basis functions at ``inputs``.
+
+        :param inputs: The evaluation points, a tensor with the shape ``[N, D]``.
+
+        :return: A tensor with the shape ``[N, 2M^D]``.
+        """
+        const = tf.tile(tf.sqrt(self.kernel.variance * self.weights), multiples=[2])
+        bases = _mapping_concat(inputs, self.abscissa, lengthscales=self.kernel.lengthscales)
+        output = const * bases
+        tf.ensure_shape(output, self.compute_output_shape(inputs.shape))
+        return output
diff --git a/gpflux/layers/basis_functions/random_fourier_features.py b/gpflux/layers/basis_functions/fourier_features/random.py
similarity index 83%
rename from gpflux/layers/basis_functions/random_fourier_features.py
rename to gpflux/layers/basis_functions/fourier_features/random.py
index bba25995..fe478d1d 100644
--- a/gpflux/layers/basis_functions/random_fourier_features.py
+++ b/gpflux/layers/basis_functions/fourier_features/random.py
@@ -23,7 +23,7 @@
 import gpflow
 from gpflow.base import DType, TensorType
 
-from gpflux.layers.basis_functions.utils import (
+from gpflux.layers.basis_functions.fourier_features.utils import (
     RFF_SUPPORTED_KERNELS,
     _mapping_concat,
     _mapping_cosine,
@@ -32,16 +32,9 @@
 )
 from gpflux.types import ShapeType
 
-"""
-Kernels supported by :class:`RandomFourierFeatures`.
-
-You can build RFF for shift-invariant stationary kernels from which you can
-sample frequencies from their power spectrum, following Bochner's theorem.
-"""
-
 
 class RandomFourierFeaturesBase(tf.keras.layers.Layer):
-    def __init__(self, kernel: gpflow.kernels.Kernel, output_dim: int, **kwargs: Mapping):
+    def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
         """
         :param kernel: kernel to approximate using a set of random features.
         :param output_dim: total number of basis functions used to approximate
@@ -51,13 +44,24 @@ def __init__(self, kernel: gpflow.kernels.Kernel, output_dim: int, **kwargs: Map
 
         assert isinstance(kernel, RFF_SUPPORTED_KERNELS), "Unsupported Kernel"
         self.kernel = kernel
-        self.output_dim = output_dim  # M
+        self.n_components = n_components  # M: number of Monte Carlo samples
         if kwargs.get("input_dim", None):
             self._input_dim = kwargs["input_dim"]
             self.build(tf.TensorShape([self._input_dim]))
         else:
             self._input_dim = None
 
+    def build(self, input_shape: ShapeType) -> None:
+        """
+        Creates the variables of the layer.
+        See `tf.keras.layers.Layer.build()
+        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#build>`_.
+        """
+        input_dim = input_shape[-1]
+        self._weights_build(input_dim, n_components=self.n_components)
+
+        super().build(input_shape)
+
     def _weights_build(self, input_dim: int, n_components: int) -> None:
         shape = (n_components, input_dim)
         self.W = self.add_weight(
@@ -85,7 +89,8 @@ def compute_output_shape(self, input_shape: ShapeType) -> tf.TensorShape:
         # TODO: Keras docs say "If the layer has not been built, this method
         # will call `build` on the layer." -- do we need to do so?
         tensor_shape = tf.TensorShape(input_shape).with_rank(2)
-        return tensor_shape[:-1].concatenate(self.output_dim)
+        output_dim = self.n_components
+        return tensor_shape[:-1].concatenate(output_dim)
 
     def get_config(self) -> Mapping:
         """
@@ -95,7 +100,7 @@ def get_config(self) -> Mapping:
         """
         config = super().get_config()
         config.update(
-            {"kernel": self.kernel, "output_dim": self.output_dim, "input_dim": self._input_dim}
+            {"kernel": self.kernel, "n_components": self.n_components, "input_dim": self._input_dim}
         )
 
         return config
@@ -135,25 +140,17 @@ class RandomFourierFeatures(RandomFourierFeaturesBase):
     counterpart :class:`RandomFourierFeaturesCosine` :cite:p:`sutherland2015error`.
     """
 
-    def __init__(self, kernel: gpflow.kernels.Kernel, output_dim: int, **kwargs: Mapping):
-        """
-        :param kernel: kernel to approximate using a set of random features.
-        :param output_dim: total number of basis functions used to approximate
-            the kernel.
-        """
-        assert not output_dim % 2, "must specify an even number of random features"
-        super().__init__(kernel, output_dim, **kwargs)
-
-    def build(self, input_shape: ShapeType) -> None:
+    def compute_output_shape(self, input_shape: ShapeType) -> tf.TensorShape:
         """
-        Creates the variables of the layer.
-        See `tf.keras.layers.Layer.build()
-        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#build>`_.
+        Computes the output shape of the layer.
+        See `tf.keras.layers.Layer.compute_output_shape()
+        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#compute_output_shape>`_.
         """
-        input_dim = input_shape[-1]
-        self._weights_build(input_dim, n_components=self.output_dim // 2)
-
-        super().build(input_shape)
+        # TODO: Keras docs say "If the layer has not been built, this method
+        # will call `build` on the layer." -- do we need to do so?
+        tensor_shape = tf.TensorShape(input_shape).with_rank(2)
+        output_dim = 2 * self.n_components
+        return tensor_shape[:-1].concatenate(output_dim)
 
     def call(self, inputs: TensorType) -> tf.Tensor:
         """
@@ -161,15 +158,11 @@ def call(self, inputs: TensorType) -> tf.Tensor:
 
         :param inputs: The evaluation points, a tensor with the shape ``[N, D]``.
 
-        :return: A tensor with the shape ``[N, M]``.
+        :return: A tensor with the shape ``[N, 2M]``.
         """
-        output = _mapping_concat(
-            inputs,
-            self.W,
-            variance=self.kernel.variance,
-            lengthscales=self.kernel.lengthscales,
-            n_components=self.output_dim,
-        )
+        constant = tf.sqrt(2.0 * self.kernel.variance / self.n_components)
+        bases = _mapping_concat(inputs, self.W, lengthscales=self.kernel.lengthscales)
+        output = constant * bases
         tf.ensure_shape(output, self.compute_output_shape(inputs.shape))
         return output
 
@@ -237,13 +230,8 @@ def call(self, inputs: TensorType) -> tf.Tensor:
 
         :return: A tensor with the shape ``[N, M]``.
         """
-        output = _mapping_cosine(
-            inputs,
-            self.W,
-            self.b,
-            variance=self.kernel.variance,
-            lengthscales=self.kernel.lengthscales,
-            n_components=self.output_dim,
-        )
+        constant = tf.sqrt(2.0 * self.kernel.variance / self.n_components)
+        bases = _mapping_cosine(inputs, self.W, lengthscales=self.kernel.lengthscales)
+        output = constant * bases
         tf.ensure_shape(output, self.compute_output_shape(inputs.shape))
         return output
diff --git a/gpflux/layers/basis_functions/utils.py b/gpflux/layers/basis_functions/fourier_features/utils.py
similarity index 83%
rename from gpflux/layers/basis_functions/utils.py
rename to gpflux/layers/basis_functions/fourier_features/utils.py
index 09058436..2b93b3ad 100644
--- a/gpflux/layers/basis_functions/utils.py
+++ b/gpflux/layers/basis_functions/fourier_features/utils.py
@@ -25,6 +25,16 @@
 
 from gpflux.types import ShapeType
 
+QFF_SUPPORTED_KERNELS: Tuple[Type[gpflow.kernels.Stationary], ...] = (
+    gpflow.kernels.SquaredExponential,
+)
+
+"""
+Kernels supported by :class:`RandomFourierFeatures`.
+
+You can build RFF for shift-invariant stationary kernels from which you can
+sample frequencies from their power spectrum, following Bochner's theorem.
+"""
 RFF_SUPPORTED_KERNELS: Tuple[Type[gpflow.kernels.Stationary], ...] = (
     gpflow.kernels.SquaredExponential,
     gpflow.kernels.Matern12,
@@ -78,36 +88,36 @@ def _mapping_cosine(
     X: TensorType,
     W: TensorType,
     b: TensorType,
-    variance: TensorType,
+    # variance: TensorType,
     lengthscales: TensorType,
-    n_components: int,
+    # n_components: int,
 ) -> TensorType:
     """
     Feature map for random Fourier features (RFF) as originally prescribed
     by Rahimi & Recht, 2007 :cite:p:`rahimi2007random`.
     See also :cite:p:`sutherland2015error` for additional details.
     """
-    constant = tf.sqrt(2.0 * variance / n_components)
+    # constant = tf.sqrt(2.0 * variance / n_components)
     X_scaled = tf.divide(X, lengthscales)  # [N, D]
     proj = tf.matmul(X_scaled, W, transpose_b=True)  # [N, M]
     bases = tf.cos(proj + b)  # [N, M]
-    return constant * bases  # [N, M]
+    return bases
 
 
 def _mapping_concat(
     X: TensorType,
     W: TensorType,
-    variance: TensorType,
+    # variance: TensorType,
     lengthscales: TensorType,
-    n_components: int,
+    # n_components: int,
 ) -> TensorType:
     """
     Feature map for random Fourier features (RFF) as originally prescribed
     by Rahimi & Recht, 2007 :cite:p:`rahimi2007random`.
     See also :cite:p:`sutherland2015error` for additional details.
     """
-    constant = tf.sqrt(2.0 * variance / n_components)
+    # constant = tf.sqrt(2.0 * variance / n_components)
     X_scaled = tf.divide(X, lengthscales)  # [N, D]
-    proj = tf.matmul(X_scaled, W, transpose_b=True)  # [N, M // 2]
-    bases = tf.concat([tf.sin(proj), tf.cos(proj)], axis=-1)  # [N, M]
-    return constant * bases  # [N, M]
+    proj = tf.matmul(X_scaled, W, transpose_b=True)  # [N, M]
+    bases = tf.concat([tf.sin(proj), tf.cos(proj)], axis=-1)  # [N, 2M]
+    return bases
diff --git a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py b/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
index d1c6cf85..7e42ecbf 100644
--- a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
+++ b/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
@@ -21,11 +21,11 @@
 
 import gpflow
 
-from gpflux.layers.basis_functions.random_fourier_features import (
-    RFF_SUPPORTED_KERNELS,
+from gpflux.layers.basis_functions.fourier_features import (
     RandomFourierFeatures,
     RandomFourierFeaturesCosine,
 )
+from gpflux.layers.basis_functions.fourier_features.utils import RFF_SUPPORTED_KERNELS
 
 
 @pytest.fixture(name="n_dims", params=[1, 2, 3, 5, 10, 20])
diff --git a/tests/gpflux/sampling/test_sample.py b/tests/gpflux/sampling/test_sample.py
index 50096d01..8255b691 100644
--- a/tests/gpflux/sampling/test_sample.py
+++ b/tests/gpflux/sampling/test_sample.py
@@ -20,7 +20,7 @@
 import gpflow
 from gpflow.config import default_float, default_jitter
 
-from gpflux.layers.basis_functions.random_fourier_features import RandomFourierFeatures
+from gpflux.layers.basis_functions.fourier_features import RandomFourierFeatures
 from gpflux.sampling.kernel_with_feature_decomposition import KernelWithFeatureDecomposition
 from gpflux.sampling.sample import Sample, efficient_sample
 

From 92e168dab48a4c4a3c9f246af95dfb6f98cd8374 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Mon, 11 Oct 2021 13:30:16 +0100
Subject: [PATCH 02/23] removed old code

---
 .../layers/basis_functions/fourier_features/random.py  |  2 --
 .../layers/basis_functions/fourier_features/utils.py   | 10 ++--------
 2 files changed, 2 insertions(+), 10 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/random.py b/gpflux/layers/basis_functions/fourier_features/random.py
index fe478d1d..a49762f2 100644
--- a/gpflux/layers/basis_functions/fourier_features/random.py
+++ b/gpflux/layers/basis_functions/fourier_features/random.py
@@ -59,7 +59,6 @@ def build(self, input_shape: ShapeType) -> None:
         """
         input_dim = input_shape[-1]
         self._weights_build(input_dim, n_components=self.n_components)
-
         super().build(input_shape)
 
     def _weights_build(self, input_dim: int, n_components: int) -> None:
@@ -206,7 +205,6 @@ def build(self, input_shape: ShapeType) -> None:
         input_dim = input_shape[-1]
         self._weights_build(input_dim, n_components=self.output_dim)
         self._bias_build(n_components=self.output_dim)
-
         super().build(input_shape)
 
     def _bias_build(self, n_components: int) -> None:
diff --git a/gpflux/layers/basis_functions/fourier_features/utils.py b/gpflux/layers/basis_functions/fourier_features/utils.py
index 2b93b3ad..8d99dc29 100644
--- a/gpflux/layers/basis_functions/fourier_features/utils.py
+++ b/gpflux/layers/basis_functions/fourier_features/utils.py
@@ -88,35 +88,29 @@ def _mapping_cosine(
     X: TensorType,
     W: TensorType,
     b: TensorType,
-    # variance: TensorType,
     lengthscales: TensorType,
-    # n_components: int,
 ) -> TensorType:
     """
     Feature map for random Fourier features (RFF) as originally prescribed
     by Rahimi & Recht, 2007 :cite:p:`rahimi2007random`.
     See also :cite:p:`sutherland2015error` for additional details.
     """
-    # constant = tf.sqrt(2.0 * variance / n_components)
     X_scaled = tf.divide(X, lengthscales)  # [N, D]
-    proj = tf.matmul(X_scaled, W, transpose_b=True)  # [N, M]
-    bases = tf.cos(proj + b)  # [N, M]
+    proj = tf.matmul(X_scaled, W, transpose_b=True) + b  # [N, M]
+    bases = tf.cos(proj)  # [N, M]
     return bases
 
 
 def _mapping_concat(
     X: TensorType,
     W: TensorType,
-    # variance: TensorType,
     lengthscales: TensorType,
-    # n_components: int,
 ) -> TensorType:
     """
     Feature map for random Fourier features (RFF) as originally prescribed
     by Rahimi & Recht, 2007 :cite:p:`rahimi2007random`.
     See also :cite:p:`sutherland2015error` for additional details.
     """
-    # constant = tf.sqrt(2.0 * variance / n_components)
     X_scaled = tf.divide(X, lengthscales)  # [N, D]
     proj = tf.matmul(X_scaled, W, transpose_b=True)  # [N, M]
     bases = tf.concat([tf.sin(proj), tf.cos(proj)], axis=-1)  # [N, 2M]

From e8ad3890d25cc2f176c47b1feabc527249942424 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Mon, 11 Oct 2021 13:31:48 +0100
Subject: [PATCH 03/23] minor fix

---
 gpflux/layers/basis_functions/fourier_features/random.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/random.py b/gpflux/layers/basis_functions/fourier_features/random.py
index a49762f2..88ca3de6 100644
--- a/gpflux/layers/basis_functions/fourier_features/random.py
+++ b/gpflux/layers/basis_functions/fourier_features/random.py
@@ -203,7 +203,6 @@ def build(self, input_shape: ShapeType) -> None:
         <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#build>`_.
         """
         input_dim = input_shape[-1]
-        self._weights_build(input_dim, n_components=self.output_dim)
         self._bias_build(n_components=self.output_dim)
         super().build(input_shape)
 

From ddfa6b73ca5f7f5c35f924ad0551b6e584f07081 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Mon, 11 Oct 2021 14:33:30 +0100
Subject: [PATCH 04/23] added submodule for abstract base class

---
 .../notebooks/efficient_posterior_sampling.py |  8 +-
 docs/notebooks/efficient_sampling.py          |  4 +-
 docs/notebooks/weight_space_approximation.py  | 12 +--
 .../fourier_features/__init__.py              |  2 +
 .../basis_functions/fourier_features/base.py  | 83 ++++++++++++++++
 .../fourier_features/quadrature.py            | 64 +++---------
 .../fourier_features/random.py                | 98 ++++---------------
 .../basis_functions/fourier_features/utils.py |  4 +-
 .../test_random_fourier_features.py           | 19 +---
 tests/gpflux/sampling/test_sample.py          |  4 +-
 10 files changed, 136 insertions(+), 162 deletions(-)
 create mode 100644 gpflux/layers/basis_functions/fourier_features/base.py

diff --git a/docs/notebooks/efficient_posterior_sampling.py b/docs/notebooks/efficient_posterior_sampling.py
index 49937e90..cb9a561b 100644
--- a/docs/notebooks/efficient_posterior_sampling.py
+++ b/docs/notebooks/efficient_posterior_sampling.py
@@ -62,7 +62,7 @@
 
 that combines both feature approximations and exact kernel evaluations from the Matheron function and weight space approximation formulas above.
 
-The subsequent experiments demonstrate the qualitative efficiency of the hybrid rule when compared to the vanilla Matheron weight space approximation, in terms of the Wasserstein distance to the exact posterior GP. To conduct these experiments, the required classes in `gpflux` are `RandomFourierFeatures`, to approximate a stationary kernel with finitely many random Fourier features $\phi_d(\cdot)$ according to Bochner's theorem and following Rahimi and Recht "Random features for large-scale kernel machines" (NeurIPS, 2007), and `KernelWithFeatureDecomposition`, to approximate a kernel with a specified set of feature functions.
+The subsequent experiments demonstrate the qualitative efficiency of the hybrid rule when compared to the vanilla Matheron weight space approximation, in terms of the Wasserstein distance to the exact posterior GP. To conduct these experiments, the required classes in `gpflux` are `RandomFourierFeaturesCosine`, to approximate a stationary kernel with finitely many random Fourier features $\phi_d(\cdot)$ according to Bochner's theorem and following Rahimi and Recht "Random features for large-scale kernel machines" (NeurIPS, 2007), and `KernelWithFeatureDecomposition`, to approximate a kernel with a specified set of feature functions.
 """
 
 # %%
@@ -79,7 +79,7 @@
 from gpflow.kernels import RBF, Matern52
 from gpflow.models import GPR
 
-from gpflux.layers.basis_functions.random_fourier_features import RandomFourierFeatures
+from gpflux.layers.basis_functions.fourier_features import RandomFourierFeaturesCosine
 from gpflux.sampling.kernel_with_feature_decomposition import KernelWithFeatureDecomposition
 
 # %% [markdown]
@@ -253,9 +253,9 @@ def conduct_experiment(num_input_dimensions, num_train_samples, num_features):
     exact_kernel = kernel_class(lengthscales=lengthscale)
 
     # weight space approximated kernel
-    feature_functions = RandomFourierFeatures(
+    feature_functions = RandomFourierFeaturesCosine(
         kernel=kernel_class(lengthscales=lengthscale),
-        output_dim=num_features,
+        n_components=num_features,
         dtype=default_float(),
     )
     feature_coefficients = np.ones((num_features, 1), dtype=default_float())
diff --git a/docs/notebooks/efficient_sampling.py b/docs/notebooks/efficient_sampling.py
index 79420344..990a1cf4 100644
--- a/docs/notebooks/efficient_sampling.py
+++ b/docs/notebooks/efficient_sampling.py
@@ -36,7 +36,7 @@
 
 from gpflow.config import default_float
 
-from gpflux.layers.basis_functions.random_fourier_features import RandomFourierFeatures
+from gpflux.layers.basis_functions.fourier_features import RandomFourierFeaturesCosine
 from gpflux.sampling import KernelWithFeatureDecomposition
 from gpflux.models.deep_gp import sample_dgp
 
@@ -71,7 +71,7 @@
 gpflow.utilities.set_trainable(inducing_variable, False)
 
 num_rff = 1000
-eigenfunctions = RandomFourierFeatures(kernel, num_rff, dtype=default_float())
+eigenfunctions = RandomFourierFeaturesCosine(kernel, num_rff, dtype=default_float())
 eigenvalues = np.ones((num_rff, 1), dtype=default_float())
 kernel_with_features = KernelWithFeatureDecomposition(kernel, eigenfunctions, eigenvalues)
 
diff --git a/docs/notebooks/weight_space_approximation.py b/docs/notebooks/weight_space_approximation.py
index 0f7a5bc2..ef4b4fe4 100644
--- a/docs/notebooks/weight_space_approximation.py
+++ b/docs/notebooks/weight_space_approximation.py
@@ -42,7 +42,7 @@
 
 The advantage of expressing a Gaussian process in weight space is that functions are represented as weight vectors $\textbf{w}$ (rather than actual functions $f(\cdot)$) from which samples can be obtained a priori without knowing where the function should be evaluated. When expressing a Gaussian process in function space view the latter is not possible, i.e. a function $f(\cdot)$ cannot be sampled without knowing where to evaluate the function, namely at $\{X_{n^\star}^\star\}_{n^\star=1,...,N^\star}$. Weight space approximated Gaussian processes therefore hold the potential to sample efficiently from Gaussian process posteriors, which is desirable in vanilla supervised learning but also in domains such as Bayesian optimisation or model-based reinforcement learning.
 
-In the following example, we compare a weight space approximated GPR model (WSA model) with both a proper GPR model and a sparse variational Gaussian Process model (SVGP). GPR models and SVGP models are implemented in `gpflow`, but the two necessary ingredients for building the WSA model are part of `gpflux`: these are random Fourier feature functions via the `RandomFourierFeatures` class, and approximate kernels based on Bochner's theorem (or any other theorem that approximates a kernel with a finite number of feature functions, e.g. Mercer) via the `KernelWithFeatureDecomposition` class.
+In the following example, we compare a weight space approximated GPR model (WSA model) with both a proper GPR model and a sparse variational Gaussian Process model (SVGP). GPR models and SVGP models are implemented in `gpflow`, but the two necessary ingredients for building the WSA model are part of `gpflux`: these are random Fourier feature functions via the `RandomFourierFeaturesCosine` class, and approximate kernels based on Bochner's theorem (or any other theorem that approximates a kernel with a finite number of feature functions, e.g. Mercer) via the `KernelWithFeatureDecomposition` class.
 """
 
 # %%
@@ -61,7 +61,7 @@
 from gpflow.likelihoods import Gaussian
 from gpflow.inducing_variables import InducingPoints
 
-from gpflux.layers.basis_functions.random_fourier_features import RandomFourierFeatures
+from gpflux.layers.basis_functions.fourier_features import RandomFourierFeaturesCosine
 from gpflux.sampling.kernel_with_feature_decomposition import KernelWithFeatureDecomposition
 
 # %% [markdown]
@@ -77,7 +77,7 @@
 # experiment parameters that are the same for both sets of experiments
 X_interval = [0.14, 0.5]  # interval where training points live
 lengthscale = 0.1  # lengthscale for the kernel (which is not learned in all experiments, the kernel variance is 1)
-number_of_basis_functions = 2000  # number of basis functions for weight-space approximated kernels
+number_of_features = 2000  # number of basis functions for weight-space approximated kernels
 noise_variance = 1e-3  # noise variance of the likelihood (which is not learned in all experiments)
 number_of_test_samples = 1024  # number of evaluation points for prediction
 number_of_function_samples = (
@@ -264,12 +264,12 @@ def optimize_model_with_scipy(model):
     )
 
     # create exact GPR model with weight-space approximated kernel (WSA model)
-    feature_functions = RandomFourierFeatures(
+    feature_functions = RandomFourierFeaturesCosine(
         kernel=kernel_class(lengthscales=lengthscale),
-        output_dim=number_of_basis_functions,
+        n_components=number_of_features,
         dtype=default_float(),
     )
-    feature_coefficients = np.ones((number_of_basis_functions, 1), dtype=default_float())
+    feature_coefficients = np.ones((number_of_features, 1), dtype=default_float())
     kernel = KernelWithFeatureDecomposition(
         kernel=None, feature_functions=feature_functions, feature_coefficients=feature_coefficients
     )
diff --git a/gpflux/layers/basis_functions/fourier_features/__init__.py b/gpflux/layers/basis_functions/fourier_features/__init__.py
index 989028b1..54a7daaf 100644
--- a/gpflux/layers/basis_functions/fourier_features/__init__.py
+++ b/gpflux/layers/basis_functions/fourier_features/__init__.py
@@ -23,3 +23,5 @@
     RandomFourierFeatures,
     RandomFourierFeaturesCosine,
 )
+
+__all__ = ["QuadratureFourierFeatures", "RandomFourierFeatures", "RandomFourierFeaturesCosine"]
diff --git a/gpflux/layers/basis_functions/fourier_features/base.py b/gpflux/layers/basis_functions/fourier_features/base.py
new file mode 100644
index 00000000..48df848b
--- /dev/null
+++ b/gpflux/layers/basis_functions/fourier_features/base.py
@@ -0,0 +1,83 @@
+from abc import ABC, abstractmethod
+from typing import Mapping
+
+import tensorflow as tf
+
+import gpflow
+from gpflow.base import TensorType
+
+from gpflux.types import ShapeType
+
+
+class FourierFeaturesBase(ABC, tf.keras.layers.Layer):
+    def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
+        """
+        :param kernel: kernel to approximate using a set of random features.
+        :param output_dim: total number of basis functions used to approximate
+            the kernel.
+        """
+        super().__init__(**kwargs)
+        self.kernel = kernel
+        self.n_components = n_components  # M: number of Monte Carlo samples
+        if kwargs.get("input_dim", None):
+            self._input_dim = kwargs["input_dim"]
+            self.build(tf.TensorShape([self._input_dim]))
+        else:
+            self._input_dim = None
+
+    def call(self, inputs: TensorType) -> tf.Tensor:
+        """
+        Evaluate the basis functions at ``inputs``.
+
+        :param inputs: The evaluation points, a tensor with the shape ``[N, D]``.
+
+        :return: A tensor with the shape ``[N, M]``.
+        """
+        const = self._compute_constant()
+        bases = self._compute_bases(inputs)
+        output = const * bases
+        tf.ensure_shape(output, self.compute_output_shape(inputs.shape))
+        return output
+
+    def compute_output_shape(self, input_shape: ShapeType) -> tf.TensorShape:
+        """
+        Computes the output shape of the layer.
+        See `tf.keras.layers.Layer.compute_output_shape()
+        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#compute_output_shape>`_.
+        """
+        # TODO: Keras docs say "If the layer has not been built, this method
+        # will call `build` on the layer." -- do we need to do so?
+        tensor_shape = tf.TensorShape(input_shape).with_rank(2)
+        output_dim = self._compute_output_dim(input_shape)
+        return tensor_shape[:-1].concatenate(output_dim)
+
+    def get_config(self) -> Mapping:
+        """
+        Returns the config of the layer.
+        See `tf.keras.layers.Layer.get_config()
+        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#get_config>`_.
+        """
+        config = super().get_config()
+        config.update(
+            {"kernel": self.kernel, "n_components": self.n_components, "input_dim": self._input_dim}
+        )
+
+        return config
+
+    @abstractmethod
+    def _compute_output_dim(self, input_shape: ShapeType) -> int:
+        pass
+
+    @abstractmethod
+    def _compute_constant(self) -> tf.Tensor:
+        """
+        Compute normalizing constant for basis functions.
+        """
+        pass
+
+    @abstractmethod
+    def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
+        """
+        Compute basis functions.
+        """
+        pass
diff --git a/gpflux/layers/basis_functions/fourier_features/quadrature.py b/gpflux/layers/basis_functions/fourier_features/quadrature.py
index 6a4f4782..94892f60 100644
--- a/gpflux/layers/basis_functions/fourier_features/quadrature.py
+++ b/gpflux/layers/basis_functions/fourier_features/quadrature.py
@@ -23,30 +23,18 @@
 from gpflow.base import TensorType
 from gpflow.quadrature.gauss_hermite import ndgh_points_and_weights
 
+from gpflux.layers.basis_functions.fourier_features.base import FourierFeaturesBase
 from gpflux.layers.basis_functions.fourier_features.utils import (
     QFF_SUPPORTED_KERNELS,
-    _mapping_concat,
+    _bases_concat,
 )
 from gpflux.types import ShapeType
 
 
-class QuadratureFourierFeatures(tf.keras.layers.Layer):
+class QuadratureFourierFeatures(FourierFeaturesBase):
     def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
-        """
-        :param kernel: kernel to approximate using a set of random features.
-        :param output_dim: total number of basis functions used to approximate
-            the kernel.
-        """
-        super().__init__(**kwargs)
-
         assert isinstance(kernel, QFF_SUPPORTED_KERNELS), "Unsupported Kernel"
-        self.kernel = kernel
-        self.n_components = n_components  # M: number of quadrature points
-        if kwargs.get("input_dim", None):
-            self._input_dim = kwargs["input_dim"]
-            self.build(tf.TensorShape([self._input_dim]))
-        else:
-            self._input_dim = None
+        super(QuadratureFourierFeatures, self).__init__(kernel, n_components, **kwargs)
 
     def build(self, input_shape: ShapeType) -> None:
         """
@@ -63,45 +51,15 @@ def build(self, input_shape: ShapeType) -> None:
         self.abscissa = tf.Variable(initial_value=abscissa_value, trainable=False)
         # Gauss-Hermite weights (not to be confused with random Fourier feature
         # weights or NN weights)
-        self.weights = tf.Variable(initiial_value=weights_value, trainable=False)
+        self.weights = tf.Variable(initial_value=weights_value, trainable=False)
         super().build(input_shape)
 
-    def compute_output_shape(self, input_shape: ShapeType) -> tf.TensorShape:
-        """
-        Computes the output shape of the layer.
-        See `tf.keras.layers.Layer.compute_output_shape()
-        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#compute_output_shape>`_.
-        """
-        # TODO: Keras docs say "If the layer has not been built, this method
-        # will call `build` on the layer." -- do we need to do so?
+    def _compute_output_dim(self, input_shape: ShapeType) -> int:
         input_dim = input_shape[-1]
-        output_dim = 2 * self.n_components ** input_dim
-        tensor_shape = tf.TensorShape(input_shape).with_rank(2)
-        return tensor_shape[:-1].concatenate(output_dim)
+        return 2 * self.n_components ** input_dim
 
-    def get_config(self) -> Mapping:
-        """
-        Returns the config of the layer.
-        See `tf.keras.layers.Layer.get_config()
-        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#get_config>`_.
-        """
-        config = super().get_config()
-        config.update(
-            {"kernel": self.kernel, "n_components": self.n_components, "input_dim": self._input_dim}
-        )
-
-        return config
-
-    def call(self, inputs: TensorType) -> tf.Tensor:
-        """
-        Evaluate the basis functions at ``inputs``.
-
-        :param inputs: The evaluation points, a tensor with the shape ``[N, D]``.
+    def _compute_constant(self) -> tf.Tensor:
+        return tf.tile(tf.sqrt(self.kernel.variance * self.weights), multiples=[2])
 
-        :return: A tensor with the shape ``[N, 2M^D]``.
-        """
-        const = tf.tile(tf.sqrt(self.kernel.variance * self.weights), multiples=[2])
-        bases = _mapping_concat(inputs, self.abscissa, lengthscales=self.kernel.lengthscales)
-        output = const * bases
-        tf.ensure_shape(output, self.compute_output_shape(inputs.shape))
-        return output
+    def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
+        return _bases_concat(inputs, self.abscissa, lengthscales=self.kernel.lengthscales)
diff --git a/gpflux/layers/basis_functions/fourier_features/random.py b/gpflux/layers/basis_functions/fourier_features/random.py
index 88ca3de6..e1232ec9 100644
--- a/gpflux/layers/basis_functions/fourier_features/random.py
+++ b/gpflux/layers/basis_functions/fourier_features/random.py
@@ -23,33 +23,21 @@
 import gpflow
 from gpflow.base import DType, TensorType
 
+from gpflux.layers.basis_functions.fourier_features.base import FourierFeaturesBase
 from gpflux.layers.basis_functions.fourier_features.utils import (
     RFF_SUPPORTED_KERNELS,
-    _mapping_concat,
-    _mapping_cosine,
+    _bases_concat,
+    _bases_cosine,
     _matern_number,
     _sample_students_t,
 )
 from gpflux.types import ShapeType
 
 
-class RandomFourierFeaturesBase(tf.keras.layers.Layer):
+class RandomFourierFeaturesBase(FourierFeaturesBase):
     def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
-        """
-        :param kernel: kernel to approximate using a set of random features.
-        :param output_dim: total number of basis functions used to approximate
-            the kernel.
-        """
-        super().__init__(**kwargs)
-
         assert isinstance(kernel, RFF_SUPPORTED_KERNELS), "Unsupported Kernel"
-        self.kernel = kernel
-        self.n_components = n_components  # M: number of Monte Carlo samples
-        if kwargs.get("input_dim", None):
-            self._input_dim = kwargs["input_dim"]
-            self.build(tf.TensorShape([self._input_dim]))
-        else:
-            self._input_dim = None
+        super(RandomFourierFeaturesBase, self).__init__(kernel, n_components, **kwargs)
 
     def build(self, input_shape: ShapeType) -> None:
         """
@@ -59,7 +47,7 @@ def build(self, input_shape: ShapeType) -> None:
         """
         input_dim = input_shape[-1]
         self._weights_build(input_dim, n_components=self.n_components)
-        super().build(input_shape)
+        super(RandomFourierFeaturesBase, self).build(input_shape)
 
     def _weights_build(self, input_dim: int, n_components: int) -> None:
         shape = (n_components, input_dim)
@@ -79,30 +67,11 @@ def _weights_init(self, shape: TensorType, dtype: Optional[DType] = None) -> Ten
             nu = 2.0 * p + 1.0  # degrees of freedom
             return _sample_students_t(nu, shape, dtype)
 
-    def compute_output_shape(self, input_shape: ShapeType) -> tf.TensorShape:
-        """
-        Computes the output shape of the layer.
-        See `tf.keras.layers.Layer.compute_output_shape()
-        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#compute_output_shape>`_.
-        """
-        # TODO: Keras docs say "If the layer has not been built, this method
-        # will call `build` on the layer." -- do we need to do so?
-        tensor_shape = tf.TensorShape(input_shape).with_rank(2)
-        output_dim = self.n_components
-        return tensor_shape[:-1].concatenate(output_dim)
-
-    def get_config(self) -> Mapping:
+    def _compute_constant(self) -> tf.Tensor:
         """
-        Returns the config of the layer.
-        See `tf.keras.layers.Layer.get_config()
-        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#get_config>`_.
+        Compute normalizing constant for basis functions.
         """
-        config = super().get_config()
-        config.update(
-            {"kernel": self.kernel, "n_components": self.n_components, "input_dim": self._input_dim}
-        )
-
-        return config
+        return tf.sqrt(2.0 * self.kernel.variance / self.n_components)
 
 
 class RandomFourierFeatures(RandomFourierFeaturesBase):
@@ -139,31 +108,11 @@ class RandomFourierFeatures(RandomFourierFeaturesBase):
     counterpart :class:`RandomFourierFeaturesCosine` :cite:p:`sutherland2015error`.
     """
 
-    def compute_output_shape(self, input_shape: ShapeType) -> tf.TensorShape:
-        """
-        Computes the output shape of the layer.
-        See `tf.keras.layers.Layer.compute_output_shape()
-        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#compute_output_shape>`_.
-        """
-        # TODO: Keras docs say "If the layer has not been built, this method
-        # will call `build` on the layer." -- do we need to do so?
-        tensor_shape = tf.TensorShape(input_shape).with_rank(2)
-        output_dim = 2 * self.n_components
-        return tensor_shape[:-1].concatenate(output_dim)
-
-    def call(self, inputs: TensorType) -> tf.Tensor:
-        """
-        Evaluate the basis functions at ``inputs``.
+    def _compute_output_dim(self, input_shape: ShapeType) -> int:
+        return 2 * self.n_components
 
-        :param inputs: The evaluation points, a tensor with the shape ``[N, D]``.
-
-        :return: A tensor with the shape ``[N, 2M]``.
-        """
-        constant = tf.sqrt(2.0 * self.kernel.variance / self.n_components)
-        bases = _mapping_concat(inputs, self.W, lengthscales=self.kernel.lengthscales)
-        output = constant * bases
-        tf.ensure_shape(output, self.compute_output_shape(inputs.shape))
-        return output
+    def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
+        return _bases_concat(inputs, self.W, lengthscales=self.kernel.lengthscales)
 
 
 class RandomFourierFeaturesCosine(RandomFourierFeaturesBase):
@@ -202,9 +151,8 @@ def build(self, input_shape: ShapeType) -> None:
         See `tf.keras.layers.Layer.build()
         <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#build>`_.
         """
-        input_dim = input_shape[-1]
-        self._bias_build(n_components=self.output_dim)
-        super().build(input_shape)
+        self._bias_build(n_components=self.n_components)
+        super(RandomFourierFeaturesCosine, self).build(input_shape)
 
     def _bias_build(self, n_components: int) -> None:
         shape = (1, n_components)
@@ -219,16 +167,8 @@ def _bias_build(self, n_components: int) -> None:
     def _bias_init(self, shape: TensorType, dtype: Optional[DType] = None) -> TensorType:
         return tf.random.uniform(shape=shape, maxval=2.0 * np.pi, dtype=dtype)
 
-    def call(self, inputs: TensorType) -> tf.Tensor:
-        """
-        Evaluate the basis functions at ``inputs``.
-
-        :param inputs: The evaluation points, a tensor with the shape ``[N, D]``.
+    def _compute_output_dim(self, input_shape: ShapeType) -> int:
+        return self.n_components
 
-        :return: A tensor with the shape ``[N, M]``.
-        """
-        constant = tf.sqrt(2.0 * self.kernel.variance / self.n_components)
-        bases = _mapping_cosine(inputs, self.W, lengthscales=self.kernel.lengthscales)
-        output = constant * bases
-        tf.ensure_shape(output, self.compute_output_shape(inputs.shape))
-        return output
+    def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
+        return _bases_cosine(inputs, self.W, self.b, lengthscales=self.kernel.lengthscales)
diff --git a/gpflux/layers/basis_functions/fourier_features/utils.py b/gpflux/layers/basis_functions/fourier_features/utils.py
index 8d99dc29..f69bc8a5 100644
--- a/gpflux/layers/basis_functions/fourier_features/utils.py
+++ b/gpflux/layers/basis_functions/fourier_features/utils.py
@@ -84,7 +84,7 @@ def _sample_students_t(nu: float, shape: ShapeType, dtype: DType) -> TensorType:
     return students_t_rvs
 
 
-def _mapping_cosine(
+def _bases_cosine(
     X: TensorType,
     W: TensorType,
     b: TensorType,
@@ -101,7 +101,7 @@ def _mapping_cosine(
     return bases
 
 
-def _mapping_concat(
+def _bases_concat(
     X: TensorType,
     W: TensorType,
     lengthscales: TensorType,
diff --git a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py b/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
index 7e42ecbf..8eb0466b 100644
--- a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
+++ b/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
@@ -43,7 +43,7 @@ def _batch_size_fixture(request):
     return request.param
 
 
-@pytest.fixture(name="n_features", params=[2, 4, 16, 128])
+@pytest.fixture(name="n_features", params=[1, 2, 4, 20, 100])
 def _n_features_fixture(request):
     return request.param
 
@@ -53,24 +53,15 @@ def _kernel_class_fixture(request):
     return request.param
 
 
-@pytest.fixture(
-    name="random_feature_class", params=[RandomFourierFeatures, RandomFourierFeaturesCosine]
-)
+@pytest.fixture(name="random_feature_class", params=[RandomFourierFeaturesCosine])
 def _random_feature_class_fixture(request):
     return request.param
 
 
-def test_throw_for_odd_num_features():
-    kernel = gpflow.kernels.Constant()
-    with pytest.raises(AssertionError) as excinfo:
-        RandomFourierFeatures(kernel, output_dim=3)
-    assert "must specify an even number of random features" in str(excinfo.value)
-
-
 def test_throw_for_unsupported_kernel(random_feature_class):
     kernel = gpflow.kernels.Constant()
     with pytest.raises(AssertionError) as excinfo:
-        random_feature_class(kernel, output_dim=2)
+        random_feature_class(kernel, n_components=1)
     assert "Unsupported Kernel" in str(excinfo.value)
 
 
@@ -137,7 +128,7 @@ def test_keras_testing_util_layer_test_1D(kernel_class, batch_size, n_features):
         RandomFourierFeatures,
         kwargs={
             "kernel": kernel,
-            "output_dim": n_features,
+            "n_components": n_features,
             "input_dim": 1,
             "dtype": "float64",
             "dynamic": True,
@@ -155,7 +146,7 @@ def test_keras_testing_util_layer_test_multidim(kernel_class, batch_size, n_dims
         RandomFourierFeatures,
         kwargs={
             "kernel": kernel,
-            "output_dim": n_features,
+            "n_components": n_features,
             "input_dim": n_dims,
             "dtype": "float64",
             "dynamic": True,
diff --git a/tests/gpflux/sampling/test_sample.py b/tests/gpflux/sampling/test_sample.py
index 8255b691..70b27588 100644
--- a/tests/gpflux/sampling/test_sample.py
+++ b/tests/gpflux/sampling/test_sample.py
@@ -20,7 +20,7 @@
 import gpflow
 from gpflow.config import default_float, default_jitter
 
-from gpflux.layers.basis_functions.fourier_features import RandomFourierFeatures
+from gpflux.layers.basis_functions.fourier_features import RandomFourierFeaturesCosine
 from gpflux.sampling.kernel_with_feature_decomposition import KernelWithFeatureDecomposition
 from gpflux.sampling.sample import Sample, efficient_sample
 
@@ -77,7 +77,7 @@ def test_conditional_sample(kernel, inducing_variable, whiten):
 
 def test_wilson_efficient_sample(kernel, inducing_variable, whiten):
     """Smoke and consistency test for efficient sampling using Wilson"""
-    eigenfunctions = RandomFourierFeatures(kernel, 100, dtype=default_float())
+    eigenfunctions = RandomFourierFeaturesCosine(kernel, 100, dtype=default_float())
     eigenvalues = np.ones((100, 1), dtype=default_float())
     # To apply Wilson sampling we require the features and eigenvalues of the kernel
     kernel2 = KernelWithFeatureDecomposition(kernel, eigenfunctions, eigenvalues)

From a82c7f43f0b31d58c859750c2d7e8b2465d822a9 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Mon, 11 Oct 2021 15:00:01 +0100
Subject: [PATCH 05/23] rename to avoid conflict

---
 .../basis_functions/fourier_features/quadrature.py    | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/quadrature.py b/gpflux/layers/basis_functions/fourier_features/quadrature.py
index 94892f60..c5e6addd 100644
--- a/gpflux/layers/basis_functions/fourier_features/quadrature.py
+++ b/gpflux/layers/basis_functions/fourier_features/quadrature.py
@@ -43,23 +43,24 @@ def build(self, input_shape: ShapeType) -> None:
         <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#build>`_.
         """
         input_dim = input_shape[-1]
-        abscissa_value, weights_value = ndgh_points_and_weights(
+        abscissa_value, omegas_value = ndgh_points_and_weights(
             dim=input_dim, n_gh=self.n_components
         )
+        omegas_value = tf.squeeze(omegas_value, axis=-1)
 
         # Quadrature node points
-        self.abscissa = tf.Variable(initial_value=abscissa_value, trainable=False)
+        self.abscissa = tf.Variable(initial_value=abscissa_value, trainable=False)  # (M^D, D)
         # Gauss-Hermite weights (not to be confused with random Fourier feature
         # weights or NN weights)
-        self.weights = tf.Variable(initial_value=weights_value, trainable=False)
-        super().build(input_shape)
+        self.omegas = tf.Variable(initial_value=omegas_value, trainable=False)  # (M^D,)
+        super(QuadratureFourierFeatures, self).build(input_shape)
 
     def _compute_output_dim(self, input_shape: ShapeType) -> int:
         input_dim = input_shape[-1]
         return 2 * self.n_components ** input_dim
 
     def _compute_constant(self) -> tf.Tensor:
-        return tf.tile(tf.sqrt(self.kernel.variance * self.weights), multiples=[2])
+        return tf.tile(tf.sqrt(self.kernel.variance * self.omegas), multiples=[2])
 
     def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
         return _bases_concat(inputs, self.abscissa, lengthscales=self.kernel.lengthscales)

From e9d0fac85e3d2f19fcfabb679100815ab482b81b Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Mon, 11 Oct 2021 15:15:39 +0100
Subject: [PATCH 06/23] explicit is better than implicit

---
 gpflux/layers/basis_functions/fourier_features/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/base.py b/gpflux/layers/basis_functions/fourier_features/base.py
index 48df848b..5ad74992 100644
--- a/gpflux/layers/basis_functions/fourier_features/base.py
+++ b/gpflux/layers/basis_functions/fourier_features/base.py
@@ -16,7 +16,7 @@ def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: M
         :param output_dim: total number of basis functions used to approximate
             the kernel.
         """
-        super().__init__(**kwargs)
+        super(FourierFeaturesBase, self).__init__(**kwargs)
         self.kernel = kernel
         self.n_components = n_components  # M: number of Monte Carlo samples
         if kwargs.get("input_dim", None):
@@ -57,7 +57,7 @@ def get_config(self) -> Mapping:
         See `tf.keras.layers.Layer.get_config()
         <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#get_config>`_.
         """
-        config = super().get_config()
+        config = super(FourierFeaturesBase, self).get_config()
         config.update(
             {"kernel": self.kernel, "n_components": self.n_components, "input_dim": self._input_dim}
         )

From 3e448d3ee5122b0e6f9dda6414a9c5a6d4db4c36 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Mon, 11 Oct 2021 18:52:09 +0100
Subject: [PATCH 07/23] fixed critical bug with kernel amplitude

---
 .../basis_functions/fourier_features/random.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/random.py b/gpflux/layers/basis_functions/fourier_features/random.py
index e1232ec9..1a9361bc 100644
--- a/gpflux/layers/basis_functions/fourier_features/random.py
+++ b/gpflux/layers/basis_functions/fourier_features/random.py
@@ -67,12 +67,6 @@ def _weights_init(self, shape: TensorType, dtype: Optional[DType] = None) -> Ten
             nu = 2.0 * p + 1.0  # degrees of freedom
             return _sample_students_t(nu, shape, dtype)
 
-    def _compute_constant(self) -> tf.Tensor:
-        """
-        Compute normalizing constant for basis functions.
-        """
-        return tf.sqrt(2.0 * self.kernel.variance / self.n_components)
-
 
 class RandomFourierFeatures(RandomFourierFeaturesBase):
     r"""
@@ -114,6 +108,12 @@ def _compute_output_dim(self, input_shape: ShapeType) -> int:
     def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
         return _bases_concat(inputs, self.W, lengthscales=self.kernel.lengthscales)
 
+    def _compute_constant(self) -> tf.Tensor:
+        """
+        Compute normalizing constant for basis functions.
+        """
+        return tf.sqrt(self.kernel.variance / self.n_components)
+
 
 class RandomFourierFeaturesCosine(RandomFourierFeaturesBase):
     r"""
@@ -172,3 +172,9 @@ def _compute_output_dim(self, input_shape: ShapeType) -> int:
 
     def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
         return _bases_cosine(inputs, self.W, self.b, lengthscales=self.kernel.lengthscales)
+
+    def _compute_constant(self) -> tf.Tensor:
+        """
+        Compute normalizing constant for basis functions.
+        """
+        return tf.sqrt(2.0 * self.kernel.variance / self.n_components)

From 8c8e8206d78927501bf937a0033f316b05aff5ae Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Mon, 11 Oct 2021 19:44:19 +0100
Subject: [PATCH 08/23] tiny refactor

---
 .../basis_functions/fourier_features/base.py  |  3 ++-
 .../fourier_features/quadrature.py            |  2 +-
 .../fourier_features/random.py                | 15 ++++++++---
 .../basis_functions/fourier_features/utils.py | 25 +++++--------------
 4 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/base.py b/gpflux/layers/basis_functions/fourier_features/base.py
index 5ad74992..dca96974 100644
--- a/gpflux/layers/basis_functions/fourier_features/base.py
+++ b/gpflux/layers/basis_functions/fourier_features/base.py
@@ -33,8 +33,9 @@ def call(self, inputs: TensorType) -> tf.Tensor:
 
         :return: A tensor with the shape ``[N, M]``.
         """
+        X = tf.divide(inputs, self.kernel.lengthscales)  # [N, D]
         const = self._compute_constant()
-        bases = self._compute_bases(inputs)
+        bases = self._compute_bases(X)
         output = const * bases
         tf.ensure_shape(output, self.compute_output_shape(inputs.shape))
         return output
diff --git a/gpflux/layers/basis_functions/fourier_features/quadrature.py b/gpflux/layers/basis_functions/fourier_features/quadrature.py
index c5e6addd..ddae0e8d 100644
--- a/gpflux/layers/basis_functions/fourier_features/quadrature.py
+++ b/gpflux/layers/basis_functions/fourier_features/quadrature.py
@@ -63,4 +63,4 @@ def _compute_constant(self) -> tf.Tensor:
         return tf.tile(tf.sqrt(self.kernel.variance * self.omegas), multiples=[2])
 
     def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
-        return _bases_concat(inputs, self.abscissa, lengthscales=self.kernel.lengthscales)
+        return _bases_concat(inputs, self.abscissa)
diff --git a/gpflux/layers/basis_functions/fourier_features/random.py b/gpflux/layers/basis_functions/fourier_features/random.py
index 1a9361bc..20e97b17 100644
--- a/gpflux/layers/basis_functions/fourier_features/random.py
+++ b/gpflux/layers/basis_functions/fourier_features/random.py
@@ -67,6 +67,13 @@ def _weights_init(self, shape: TensorType, dtype: Optional[DType] = None) -> Ten
             nu = 2.0 * p + 1.0  # degrees of freedom
             return _sample_students_t(nu, shape, dtype)
 
+    @staticmethod
+    def rff_constant(variance: TensorType, output_dim: int) -> tf.Tensor:
+        """
+        Normalizing constant for random Fourier features.
+        """
+        return tf.sqrt(2.0 * variance / output_dim)
+
 
 class RandomFourierFeatures(RandomFourierFeaturesBase):
     r"""
@@ -106,13 +113,13 @@ def _compute_output_dim(self, input_shape: ShapeType) -> int:
         return 2 * self.n_components
 
     def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
-        return _bases_concat(inputs, self.W, lengthscales=self.kernel.lengthscales)
+        return _bases_concat(inputs, self.W)
 
     def _compute_constant(self) -> tf.Tensor:
         """
         Compute normalizing constant for basis functions.
         """
-        return tf.sqrt(self.kernel.variance / self.n_components)
+        return self.rff_constant(self.kernel.variance, 2 * self.n_components)
 
 
 class RandomFourierFeaturesCosine(RandomFourierFeaturesBase):
@@ -171,10 +178,10 @@ def _compute_output_dim(self, input_shape: ShapeType) -> int:
         return self.n_components
 
     def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
-        return _bases_cosine(inputs, self.W, self.b, lengthscales=self.kernel.lengthscales)
+        return _bases_cosine(inputs, self.W, self.b)
 
     def _compute_constant(self) -> tf.Tensor:
         """
         Compute normalizing constant for basis functions.
         """
-        return tf.sqrt(2.0 * self.kernel.variance / self.n_components)
+        return self.rff_constant(self.kernel.variance, self.n_components)
diff --git a/gpflux/layers/basis_functions/fourier_features/utils.py b/gpflux/layers/basis_functions/fourier_features/utils.py
index f69bc8a5..950f5847 100644
--- a/gpflux/layers/basis_functions/fourier_features/utils.py
+++ b/gpflux/layers/basis_functions/fourier_features/utils.py
@@ -84,34 +84,21 @@ def _sample_students_t(nu: float, shape: ShapeType, dtype: DType) -> TensorType:
     return students_t_rvs
 
 
-def _bases_cosine(
-    X: TensorType,
-    W: TensorType,
-    b: TensorType,
-    lengthscales: TensorType,
-) -> TensorType:
+def _bases_cosine(X: TensorType, W: TensorType, b: TensorType) -> TensorType:
     """
     Feature map for random Fourier features (RFF) as originally prescribed
     by Rahimi & Recht, 2007 :cite:p:`rahimi2007random`.
     See also :cite:p:`sutherland2015error` for additional details.
     """
-    X_scaled = tf.divide(X, lengthscales)  # [N, D]
-    proj = tf.matmul(X_scaled, W, transpose_b=True) + b  # [N, M]
-    bases = tf.cos(proj)  # [N, M]
-    return bases
+    proj = tf.matmul(X, W, transpose_b=True) + b  # [N, M]
+    return tf.cos(proj)  # [N, M]
 
 
-def _bases_concat(
-    X: TensorType,
-    W: TensorType,
-    lengthscales: TensorType,
-) -> TensorType:
+def _bases_concat(X: TensorType, W: TensorType) -> TensorType:
     """
     Feature map for random Fourier features (RFF) as originally prescribed
     by Rahimi & Recht, 2007 :cite:p:`rahimi2007random`.
     See also :cite:p:`sutherland2015error` for additional details.
     """
-    X_scaled = tf.divide(X, lengthscales)  # [N, D]
-    proj = tf.matmul(X_scaled, W, transpose_b=True)  # [N, M]
-    bases = tf.concat([tf.sin(proj), tf.cos(proj)], axis=-1)  # [N, 2M]
-    return bases
+    proj = tf.matmul(X, W, transpose_b=True)  # [N, M]
+    return tf.concat([tf.sin(proj), tf.cos(proj)], axis=-1)  # [N, 2M]

From 6e270ea63f76bf2d77ea45f4bc4d5fd26745c7b0 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Tue, 19 Oct 2021 10:21:12 +0100
Subject: [PATCH 09/23] updated notebook and tests

---
 .../notebooks/efficient_posterior_sampling.py | 26 +++++++++----------
 .../test_random_fourier_features.py           | 26 ++++++++++++-------
 2 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/docs/notebooks/efficient_posterior_sampling.py b/docs/notebooks/efficient_posterior_sampling.py
index cb9a561b..f4ddd00e 100644
--- a/docs/notebooks/efficient_posterior_sampling.py
+++ b/docs/notebooks/efficient_posterior_sampling.py
@@ -98,7 +98,7 @@
 num_input_dimensions = [2, 4, 8]  # number of input dimensions
 train_sample_exponents = [2, 4, 6, 8, 10]  # num_train_samples = 2 ** train_sample_exponents
 num_train_samples = [2 ** train_sample_exponent for train_sample_exponent in train_sample_exponents]
-num_features = [
+num_components = [
     1024,
     4096,
     16384,
@@ -233,7 +233,7 @@ def log10_Wasserstein_distance(
 """
 
 # %%
-def conduct_experiment(num_input_dimensions, num_train_samples, num_features):
+def conduct_experiment(num_input_dimensions, num_train_samples, num_components):
     """
     Compute the log10 Wassertein distance between the weight space approximated GP and the exact GP,
     and between the hybrid-rule approximated GP and the exact GP.
@@ -247,7 +247,7 @@ def conduct_experiment(num_input_dimensions, num_train_samples, num_features):
     lengthscale = (
         num_input_dimensions / 100.0
     ) ** 0.5  # adjust kernel lengthscale to the number of input dims
-    num_features = num_train_samples + num_features
+    num_features = num_train_samples + num_components
 
     # exact kernel
     exact_kernel = kernel_class(lengthscales=lengthscale)
@@ -255,7 +255,7 @@ def conduct_experiment(num_input_dimensions, num_train_samples, num_features):
     # weight space approximated kernel
     feature_functions = RandomFourierFeaturesCosine(
         kernel=kernel_class(lengthscales=lengthscale),
-        n_components=num_features,
+        n_components=num_components,
         dtype=default_float(),
     )
     feature_coefficients = np.ones((num_features, 1), dtype=default_float())
@@ -327,7 +327,7 @@ def conduct_experiment(num_input_dimensions, num_train_samples, num_features):
 """
 
 # %%
-def conduct_experiment_for_multiple_runs(num_input_dimensions, num_train_samples, num_features):
+def conduct_experiment_for_multiple_runs(num_input_dimensions, num_train_samples, num_components):
     """
     Conduct the experiment as specified above `num_experiment_runs` times and identify the quartiles for
     the log10 Wassertein distance between the weight space approximated GP and the exact GP,
@@ -347,7 +347,7 @@ def conduct_experiment_for_multiple_runs(num_input_dimensions, num_train_samples
         log10_ws_dist_weight, log10_ws_dist_hybrid = conduct_experiment(
             num_input_dimensions=num_input_dimensions,
             num_train_samples=num_train_samples,
-            num_features=num_features,
+            num_components=num_components,
         )
         list_of_log10_ws_dist_weight.append(log10_ws_dist_weight)
         list_of_log10_ws_dist_hybrid.append(log10_ws_dist_hybrid)
@@ -363,7 +363,7 @@ def conduct_experiment_for_multiple_runs(num_input_dimensions, num_train_samples
 """
 
 # %%
-def conduct_experiment_for_different_train_data_sizes(num_input_dimensions, num_features):
+def conduct_experiment_for_different_train_data_sizes(num_input_dimensions, num_components):
     """
     Conduct the experiment as specified above for different training dataset sizes and store the results in lists.
 
@@ -383,13 +383,13 @@ def conduct_experiment_for_different_train_data_sizes(num_input_dimensions, num_
         ) = conduct_experiment_for_multiple_runs(
             num_input_dimensions=num_input_dimensions,
             num_train_samples=nts,
-            num_features=num_features,
+            num_components=num_components,
         )
         print(
             "Completed for num input dims = "
             + str(num_input_dimensions)
             + " and feature param = "
-            + str(num_features)
+            + str(num_components)
             + " and num train samples = "
             + str(nts)
         )
@@ -420,9 +420,9 @@ def conduct_experiment_for_different_num_features(num_input_dimensions):
         []
     )  # for the analytic solution using the weight space approximated kernel
     list_of_hybrid_results = []  # for the hybrid-rule approximation
-    for nf in num_features:
+    for nf in num_components:
         weight_results, hybrid_results = conduct_experiment_for_different_train_data_sizes(
-            num_input_dimensions=num_input_dimensions, num_features=nf
+            num_input_dimensions=num_input_dimensions, num_components=nf
         )
         print()
         list_of_weight_results.append(weight_results)
@@ -454,7 +454,7 @@ def conduct_experiment_for_different_num_features(num_input_dimensions):
 
     # plot the results for the analytic solution using the weight space approximated kernel
     colors = ["bisque", "orange", "peru"]
-    assert len(colors) == len(num_features), "Number of colors must equal the number of features!"
+    assert len(colors) == len(num_components), "Number of colors must equal the number of features!"
     for j in range(len(weight_results)):
         weight_result = weight_results[j]
         axs[i].fill_between(
@@ -465,7 +465,7 @@ def conduct_experiment_for_different_num_features(num_input_dimensions):
 
     # plot the results for the hybrid-rule approximation
     colors = ["lightblue", "blue", "darkblue"]
-    assert len(colors) == len(num_features), "Number of colors must equal the number of features!"
+    assert len(colors) == len(num_components), "Number of colors must equal the number of features!"
     for j in range(len(hybrid_results)):
         hybrid_result = hybrid_results[j]
         axs[i].fill_between(
diff --git a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py b/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
index 8eb0466b..0a631b7d 100644
--- a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
+++ b/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
@@ -43,7 +43,7 @@ def _batch_size_fixture(request):
     return request.param
 
 
-@pytest.fixture(name="n_features", params=[1, 2, 4, 20, 100])
+@pytest.fixture(name="n_components", params=[1, 2, 4, 20, 100])
 def _n_features_fixture(request):
     return request.param
 
@@ -68,14 +68,14 @@ def test_throw_for_unsupported_kernel(random_feature_class):
 def test_fourier_features_can_approximate_kernel_multidim(
     random_feature_class, kernel_class, lengthscale, n_dims
 ):
-    n_features = 10000
+    n_components = 10000
     x_rows = 20
     y_rows = 30
     # ARD
     lengthscales = np.random.rand((n_dims)) * lengthscale
 
     kernel = kernel_class(lengthscales=lengthscales)
-    fourier_features = random_feature_class(kernel, n_features, dtype=tf.float64)
+    fourier_features = random_feature_class(kernel, n_components, dtype=tf.float64)
 
     x = tf.random.uniform((x_rows, n_dims), dtype=tf.float64)
     y = tf.random.uniform((y_rows, n_dims), dtype=tf.float64)
@@ -89,12 +89,18 @@ def test_fourier_features_can_approximate_kernel_multidim(
     np.testing.assert_allclose(approx_kernel_matrix, actual_kernel_matrix, atol=5e-2)
 
 
-def test_fourier_features_shapes(random_feature_class, n_features, n_dims, batch_size):
+def test_fourier_features_shapes(n_components, n_dims, batch_size):
     kernel = gpflow.kernels.SquaredExponential()
-    fourier_features = random_feature_class(kernel, n_features, dtype=tf.float64)
+    fourier_features = RandomFourierFeaturesCosine(kernel, n_components, dtype=tf.float64)
     features = fourier_features(tf.ones(shape=(batch_size, n_dims)))
+    np.testing.assert_equal(features.shape, [batch_size, n_components])
 
-    np.testing.assert_equal(features.shape, [batch_size, n_features])
+
+def test_fourier_features_shapes(n_components, n_dims, batch_size):
+    kernel = gpflow.kernels.SquaredExponential()
+    fourier_features = RandomFourierFeaturesCosine(kernel, n_components, dtype=tf.float64)
+    features = fourier_features(tf.ones(shape=(batch_size, n_dims)))
+    np.testing.assert_equal(features.shape, [batch_size, n_components])
 
 
 def test_fourier_feature_layer_compute_covariance_of_inducing_variables(
@@ -120,7 +126,7 @@ def test_fourier_feature_layer_compute_covariance_of_inducing_variables(
     np.testing.assert_allclose(approx_kernel_matrix, actual_kernel_matrix, atol=5e-2)
 
 
-def test_keras_testing_util_layer_test_1D(kernel_class, batch_size, n_features):
+def test_keras_testing_util_layer_test_1D(kernel_class, batch_size, n_components):
     kernel = kernel_class()
 
     tf.keras.utils.get_custom_objects()["RandomFourierFeatures"] = RandomFourierFeatures
@@ -128,7 +134,7 @@ def test_keras_testing_util_layer_test_1D(kernel_class, batch_size, n_features):
         RandomFourierFeatures,
         kwargs={
             "kernel": kernel,
-            "n_components": n_features,
+            "n_components": n_components,
             "input_dim": 1,
             "dtype": "float64",
             "dynamic": True,
@@ -138,7 +144,7 @@ def test_keras_testing_util_layer_test_1D(kernel_class, batch_size, n_features):
     )
 
 
-def test_keras_testing_util_layer_test_multidim(kernel_class, batch_size, n_dims, n_features):
+def test_keras_testing_util_layer_test_multidim(kernel_class, batch_size, n_dims, n_components):
     kernel = kernel_class()
 
     tf.keras.utils.get_custom_objects()["RandomFourierFeatures"] = RandomFourierFeatures
@@ -146,7 +152,7 @@ def test_keras_testing_util_layer_test_multidim(kernel_class, batch_size, n_dims
         RandomFourierFeatures,
         kwargs={
             "kernel": kernel,
-            "n_components": n_features,
+            "n_components": n_components,
             "input_dim": n_dims,
             "dtype": "float64",
             "dynamic": True,

From 871aa1ff85d2e7c35f32f92768c4500248ec283e Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Tue, 19 Oct 2021 11:24:25 +0100
Subject: [PATCH 10/23] fixed test and typos

---
 .../fourier_features/random.py                |  6 +--
 .../test_random_fourier_features.py           | 54 +++++++++----------
 2 files changed, 28 insertions(+), 32 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/random.py b/gpflux/layers/basis_functions/fourier_features/random.py
index 20e97b17..523ac9be 100644
--- a/gpflux/layers/basis_functions/fourier_features/random.py
+++ b/gpflux/layers/basis_functions/fourier_features/random.py
@@ -105,8 +105,8 @@ class RandomFourierFeatures(RandomFourierFeaturesBase):
       where :math:`p(\boldsymbol{\theta})` is the spectral density of the kernel.
 
     At least for the squared exponential kernel, this variant of the feature
-    mapping has more desirable theoretical properties than its cosine-based
-    counterpart :class:`RandomFourierFeaturesCosine` :cite:p:`sutherland2015error`.
+    mapping has more desirable theoretical properties than its counterpart form 
+    from phase-shifted cosines :class:`RandomFourierFeaturesCosine` :cite:p:`sutherland2015error`.
     """
 
     def _compute_output_dim(self, input_shape: ShapeType) -> int:
@@ -149,7 +149,7 @@ class RandomFourierFeaturesCosine(RandomFourierFeaturesBase):
       where :math:`p(\boldsymbol{\theta})` is the spectral density of the kernel
     - :math:`\tau \sim \mathcal{U}(0, 2\pi)`
 
-    Equivalent to :class:`RandomFourierFeatures` by elementary trignometric identities.
+    Equivalent to :class:`RandomFourierFeatures` by elementary trigonometric identities.
     """
 
     def build(self, input_shape: ShapeType) -> None:
diff --git a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py b/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
index 0a631b7d..84c74711 100644
--- a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
+++ b/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
@@ -48,25 +48,26 @@ def _n_features_fixture(request):
     return request.param
 
 
-@pytest.fixture(name="kernel_class", params=list(RFF_SUPPORTED_KERNELS))
-def _kernel_class_fixture(request):
+@pytest.fixture(name="kernel_cls", params=list(RFF_SUPPORTED_KERNELS))
+def _kernel_cls_fixture(request):
     return request.param
 
 
-@pytest.fixture(name="random_feature_class", params=[RandomFourierFeaturesCosine])
-def _random_feature_class_fixture(request):
+@pytest.fixture(name="random_feature_cls", params=[RandomFourierFeatures, 
+                                                   RandomFourierFeaturesCosine])
+def _random_feature_cls_fixture(request):
     return request.param
 
 
-def test_throw_for_unsupported_kernel(random_feature_class):
+def test_throw_for_unsupported_kernel(random_feature_cls):
     kernel = gpflow.kernels.Constant()
     with pytest.raises(AssertionError) as excinfo:
-        random_feature_class(kernel, n_components=1)
+        random_feature_cls(kernel, n_components=1)
     assert "Unsupported Kernel" in str(excinfo.value)
 
 
 def test_fourier_features_can_approximate_kernel_multidim(
-    random_feature_class, kernel_class, lengthscale, n_dims
+    random_feature_cls, kernel_cls, lengthscale, n_dims
 ):
     n_components = 10000
     x_rows = 20
@@ -74,8 +75,8 @@ def test_fourier_features_can_approximate_kernel_multidim(
     # ARD
     lengthscales = np.random.rand((n_dims)) * lengthscale
 
-    kernel = kernel_class(lengthscales=lengthscales)
-    fourier_features = random_feature_class(kernel, n_components, dtype=tf.float64)
+    kernel = kernel_cls(lengthscales=lengthscales)
+    fourier_features = random_feature_cls(kernel, n_components, dtype=tf.float64)
 
     x = tf.random.uniform((x_rows, n_dims), dtype=tf.float64)
     y = tf.random.uniform((y_rows, n_dims), dtype=tf.float64)
@@ -89,22 +90,8 @@ def test_fourier_features_can_approximate_kernel_multidim(
     np.testing.assert_allclose(approx_kernel_matrix, actual_kernel_matrix, atol=5e-2)
 
 
-def test_fourier_features_shapes(n_components, n_dims, batch_size):
-    kernel = gpflow.kernels.SquaredExponential()
-    fourier_features = RandomFourierFeaturesCosine(kernel, n_components, dtype=tf.float64)
-    features = fourier_features(tf.ones(shape=(batch_size, n_dims)))
-    np.testing.assert_equal(features.shape, [batch_size, n_components])
-
-
-def test_fourier_features_shapes(n_components, n_dims, batch_size):
-    kernel = gpflow.kernels.SquaredExponential()
-    fourier_features = RandomFourierFeaturesCosine(kernel, n_components, dtype=tf.float64)
-    features = fourier_features(tf.ones(shape=(batch_size, n_dims)))
-    np.testing.assert_equal(features.shape, [batch_size, n_components])
-
-
 def test_fourier_feature_layer_compute_covariance_of_inducing_variables(
-    random_feature_class, batch_size
+    random_feature_cls, batch_size
 ):
     """
     Ensure that the random fourier feature map can be used to approximate the covariance matrix
@@ -114,7 +101,7 @@ def test_fourier_feature_layer_compute_covariance_of_inducing_variables(
     n_features = 10000
 
     kernel = gpflow.kernels.SquaredExponential()
-    fourier_features = random_feature_class(kernel, n_features, dtype=tf.float64)
+    fourier_features = random_feature_cls(kernel, n_features, dtype=tf.float64)
 
     x_new = tf.ones(shape=(2 * batch_size + 1, 1), dtype=tf.float64)
 
@@ -126,8 +113,17 @@ def test_fourier_feature_layer_compute_covariance_of_inducing_variables(
     np.testing.assert_allclose(approx_kernel_matrix, actual_kernel_matrix, atol=5e-2)
 
 
-def test_keras_testing_util_layer_test_1D(kernel_class, batch_size, n_components):
-    kernel = kernel_class()
+def test_fourier_features_shapes(random_feature_cls, n_components, n_dims, batch_size):
+    input_shape = (batch_size, n_dims)
+    kernel = gpflow.kernels.SquaredExponential()
+    feature_functions = random_feature_cls(kernel, n_components, dtype=tf.float64)
+    output_shape = feature_functions.compute_output_shape(input_shape)
+    features = feature_functions(tf.ones(shape=input_shape))
+    np.testing.assert_equal(features.shape, output_shape)
+
+
+def test_keras_testing_util_layer_test_1D(kernel_cls, batch_size, n_components):
+    kernel = kernel_cls()
 
     tf.keras.utils.get_custom_objects()["RandomFourierFeatures"] = RandomFourierFeatures
     layer_test(
@@ -144,8 +140,8 @@ def test_keras_testing_util_layer_test_1D(kernel_class, batch_size, n_components
     )
 
 
-def test_keras_testing_util_layer_test_multidim(kernel_class, batch_size, n_dims, n_components):
-    kernel = kernel_class()
+def test_keras_testing_util_layer_test_multidim(kernel_cls, batch_size, n_dims, n_components):
+    kernel = kernel_cls()
 
     tf.keras.utils.get_custom_objects()["RandomFourierFeatures"] = RandomFourierFeatures
     layer_test(

From e3a37ffdfabc83a965d9312a13dff48d654783f7 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Tue, 19 Oct 2021 12:23:16 +0100
Subject: [PATCH 11/23] added tests for varying kernel amplitudes

---
 .../basis_functions/fourier_features/random.py      |  6 +++---
 .../basis_functions/test_random_fourier_features.py | 13 +++++++++----
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/random.py b/gpflux/layers/basis_functions/fourier_features/random.py
index 523ac9be..7a8fbbfa 100644
--- a/gpflux/layers/basis_functions/fourier_features/random.py
+++ b/gpflux/layers/basis_functions/fourier_features/random.py
@@ -72,7 +72,7 @@ def rff_constant(variance: TensorType, output_dim: int) -> tf.Tensor:
         """
         Normalizing constant for random Fourier features.
         """
-        return tf.sqrt(2.0 * variance / output_dim)
+        return tf.sqrt(tf.math.truediv(2.0 * variance, output_dim))
 
 
 class RandomFourierFeatures(RandomFourierFeaturesBase):
@@ -119,7 +119,7 @@ def _compute_constant(self) -> tf.Tensor:
         """
         Compute normalizing constant for basis functions.
         """
-        return self.rff_constant(self.kernel.variance, 2 * self.n_components)
+        return self.rff_constant(self.kernel.variance, output_dim=2.0 * self.n_components)
 
 
 class RandomFourierFeaturesCosine(RandomFourierFeaturesBase):
@@ -184,4 +184,4 @@ def _compute_constant(self) -> tf.Tensor:
         """
         Compute normalizing constant for basis functions.
         """
-        return self.rff_constant(self.kernel.variance, self.n_components)
+        return self.rff_constant(self.kernel.variance, output_dim=self.n_components)
diff --git a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py b/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
index 84c74711..6b7acc80 100644
--- a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
+++ b/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
@@ -33,6 +33,11 @@ def _n_dims_fixture(request):
     return request.param
 
 
+@pytest.fixture(name="variance", params=[0.5, 1.0, 2.0])
+def _variance_fixture(request):
+    return request.param
+
+
 @pytest.fixture(name="lengthscale", params=[1e-3, 0.1, 1.0, 5.0])
 def _lengthscale_fixture(request):
     return request.param
@@ -53,7 +58,7 @@ def _kernel_cls_fixture(request):
     return request.param
 
 
-@pytest.fixture(name="random_feature_cls", params=[RandomFourierFeatures, 
+@pytest.fixture(name="random_feature_cls", params=[RandomFourierFeatures,
                                                    RandomFourierFeaturesCosine])
 def _random_feature_cls_fixture(request):
     return request.param
@@ -67,15 +72,15 @@ def test_throw_for_unsupported_kernel(random_feature_cls):
 
 
 def test_fourier_features_can_approximate_kernel_multidim(
-    random_feature_cls, kernel_cls, lengthscale, n_dims
+    random_feature_cls, kernel_cls, variance, lengthscale, n_dims
 ):
-    n_components = 10000
+    n_components = 50000
     x_rows = 20
     y_rows = 30
     # ARD
     lengthscales = np.random.rand((n_dims)) * lengthscale
 
-    kernel = kernel_cls(lengthscales=lengthscales)
+    kernel = kernel_cls(variance=variance, lengthscales=lengthscales)
     fourier_features = random_feature_cls(kernel, n_components, dtype=tf.float64)
 
     x = tf.random.uniform((x_rows, n_dims), dtype=tf.float64)

From 4940edecd07adea920d9def6744fdff4874e25bb Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Tue, 19 Oct 2021 15:53:06 +0100
Subject: [PATCH 12/23] added new tests

---
 .../fourier_features/random.py                |   4 +-
 .../fourier_features/test_quadrature.py       | 133 ++++++++++++++++++
 .../test_random.py}                           |  32 ++---
 3 files changed, 151 insertions(+), 18 deletions(-)
 create mode 100644 tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
 rename tests/gpflux/layers/basis_functions/{test_random_fourier_features.py => fourier_features/test_random.py} (81%)

diff --git a/gpflux/layers/basis_functions/fourier_features/random.py b/gpflux/layers/basis_functions/fourier_features/random.py
index 7a8fbbfa..076dc374 100644
--- a/gpflux/layers/basis_functions/fourier_features/random.py
+++ b/gpflux/layers/basis_functions/fourier_features/random.py
@@ -105,7 +105,7 @@ class RandomFourierFeatures(RandomFourierFeaturesBase):
       where :math:`p(\boldsymbol{\theta})` is the spectral density of the kernel.
 
     At least for the squared exponential kernel, this variant of the feature
-    mapping has more desirable theoretical properties than its counterpart form 
+    mapping has more desirable theoretical properties than its counterpart form
     from phase-shifted cosines :class:`RandomFourierFeaturesCosine` :cite:p:`sutherland2015error`.
     """
 
@@ -119,7 +119,7 @@ def _compute_constant(self) -> tf.Tensor:
         """
         Compute normalizing constant for basis functions.
         """
-        return self.rff_constant(self.kernel.variance, output_dim=2.0 * self.n_components)
+        return self.rff_constant(self.kernel.variance, output_dim=2 * self.n_components)
 
 
 class RandomFourierFeaturesCosine(RandomFourierFeaturesBase):
diff --git a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
new file mode 100644
index 00000000..b7898d6f
--- /dev/null
+++ b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
@@ -0,0 +1,133 @@
+#
+# Copyright (c) 2021 The GPflux Contributors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import numpy as np
+import pytest
+import tensorflow as tf
+from tensorflow.python.keras.testing_utils import layer_test
+from tensorflow.python.keras.utils.kernelized_utils import inner_product
+
+import gpflow
+
+from gpflux.layers.basis_functions.fourier_features import QuadratureFourierFeatures
+from gpflux.layers.basis_functions.fourier_features.utils import QFF_SUPPORTED_KERNELS
+
+
+@pytest.fixture(name="n_dims", params=[1, 2, 3])
+def _n_dims_fixture(request):
+    return request.param
+
+
+@pytest.fixture(name="n_components", params=[1, 2, 4, 20])
+def _n_features_fixture(request):
+    return request.param
+
+
+@pytest.fixture(name="variance", params=[0.5, 1.0, 2.0])
+def _variance_fixture(request):
+    return request.param
+
+
+@pytest.fixture(name="lengthscale", params=[0.1, 1.0, 5.0])
+def _lengthscale_fixture(request):
+    return request.param
+
+
+@pytest.fixture(name="batch_size", params=[1, 10])
+def _batch_size_fixture(request):
+    return request.param
+
+
+@pytest.fixture(name="kernel_cls", params=list(QFF_SUPPORTED_KERNELS))
+def _kernel_cls_fixture(request):
+    return request.param
+
+
+def test_throw_for_unsupported_kernel():
+    kernel = gpflow.kernels.Constant()
+    with pytest.raises(AssertionError) as excinfo:
+        QuadratureFourierFeatures(kernel, n_components=1)
+    assert "Unsupported Kernel" in str(excinfo.value)
+
+
+def test_quadrature_fourier_features_can_approximate_kernel_multidim(kernel_cls, variance, lengthscale, n_dims):
+
+    n_components = 60
+
+    x_rows = 20
+    y_rows = 30
+    # ARD
+    lengthscales = np.random.rand(n_dims) * lengthscale
+
+    kernel = kernel_cls(variance=variance, lengthscales=lengthscales)
+    fourier_features = QuadratureFourierFeatures(kernel, n_components, dtype=tf.float64)
+
+    x = tf.random.uniform((x_rows, n_dims), dtype=tf.float64)
+    y = tf.random.uniform((y_rows, n_dims), dtype=tf.float64)
+
+    u = fourier_features(x)
+    v = fourier_features(y)
+    approx_kernel_matrix = inner_product(u, v)
+
+    actual_kernel_matrix = kernel.K(x, y)
+
+    np.testing.assert_allclose(approx_kernel_matrix, actual_kernel_matrix, atol=5e-2)
+
+
+def test_fourier_features_shapes(n_components, n_dims, batch_size):
+    input_shape = (batch_size, n_dims)
+    kernel = gpflow.kernels.SquaredExponential()
+    feature_functions = QuadratureFourierFeatures(kernel, n_components, dtype=tf.float64)
+    output_shape = feature_functions.compute_output_shape(input_shape)
+    output_dim = output_shape[-1]
+    assert output_dim == 2*n_components**n_dims
+    features = feature_functions(tf.ones(shape=input_shape))
+    np.testing.assert_equal(features.shape, output_shape)
+
+
+def test_keras_testing_util_layer_test_1D(kernel_cls, batch_size, n_components):
+    kernel = kernel_cls()
+
+    tf.keras.utils.get_custom_objects()["QuadratureFourierFeatures"] = QuadratureFourierFeatures
+    layer_test(
+        QuadratureFourierFeatures,
+        kwargs={
+            "kernel": kernel,
+            "n_components": n_components,
+            "input_dim": 1,
+            "dtype": "float64",
+            "dynamic": True,
+        },
+        input_shape=(batch_size, 1),
+        input_dtype="float64",
+    )
+
+
+def test_keras_testing_util_layer_test_multidim(kernel_cls, batch_size, n_dims, n_components):
+    kernel = kernel_cls()
+
+    tf.keras.utils.get_custom_objects()["QuadratureFourierFeatures"] = QuadratureFourierFeatures
+    layer_test(
+        QuadratureFourierFeatures,
+        kwargs={
+            "kernel": kernel,
+            "n_components": n_components,
+            "input_dim": n_dims,
+            "dtype": "float64",
+            "dynamic": True,
+        },
+        input_shape=(batch_size, n_dims),
+        input_dtype="float64",
+    )
diff --git a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py b/tests/gpflux/layers/basis_functions/fourier_features/test_random.py
similarity index 81%
rename from tests/gpflux/layers/basis_functions/test_random_fourier_features.py
rename to tests/gpflux/layers/basis_functions/fourier_features/test_random.py
index 6b7acc80..a7a0f7a9 100644
--- a/tests/gpflux/layers/basis_functions/test_random_fourier_features.py
+++ b/tests/gpflux/layers/basis_functions/fourier_features/test_random.py
@@ -38,7 +38,7 @@ def _variance_fixture(request):
     return request.param
 
 
-@pytest.fixture(name="lengthscale", params=[1e-3, 0.1, 1.0, 5.0])
+@pytest.fixture(name="lengthscale", params=[0.1, 1.0, 5.0])
 def _lengthscale_fixture(request):
     return request.param
 
@@ -58,30 +58,30 @@ def _kernel_cls_fixture(request):
     return request.param
 
 
-@pytest.fixture(name="random_feature_cls", params=[RandomFourierFeatures,
-                                                   RandomFourierFeaturesCosine])
-def _random_feature_cls_fixture(request):
+@pytest.fixture(name="basis_func_cls", params=[RandomFourierFeatures, RandomFourierFeaturesCosine])
+def _basis_func_cls_fixture(request):
     return request.param
 
 
-def test_throw_for_unsupported_kernel(random_feature_cls):
+def test_throw_for_unsupported_kernel(basis_func_cls):
     kernel = gpflow.kernels.Constant()
     with pytest.raises(AssertionError) as excinfo:
-        random_feature_cls(kernel, n_components=1)
+        basis_func_cls(kernel, n_components=1)
     assert "Unsupported Kernel" in str(excinfo.value)
 
 
-def test_fourier_features_can_approximate_kernel_multidim(
-    random_feature_cls, kernel_cls, variance, lengthscale, n_dims
+def test_random_fourier_features_can_approximate_kernel_multidim(
+    basis_func_cls, kernel_cls, variance, lengthscale, n_dims
 ):
-    n_components = 50000
+    n_components = 40000
+
     x_rows = 20
     y_rows = 30
     # ARD
     lengthscales = np.random.rand((n_dims)) * lengthscale
 
     kernel = kernel_cls(variance=variance, lengthscales=lengthscales)
-    fourier_features = random_feature_cls(kernel, n_components, dtype=tf.float64)
+    fourier_features = basis_func_cls(kernel, n_components, dtype=tf.float64)
 
     x = tf.random.uniform((x_rows, n_dims), dtype=tf.float64)
     y = tf.random.uniform((y_rows, n_dims), dtype=tf.float64)
@@ -95,18 +95,18 @@ def test_fourier_features_can_approximate_kernel_multidim(
     np.testing.assert_allclose(approx_kernel_matrix, actual_kernel_matrix, atol=5e-2)
 
 
-def test_fourier_feature_layer_compute_covariance_of_inducing_variables(
-    random_feature_cls, batch_size
+def test_random_fourier_feature_layer_compute_covariance_of_inducing_variables(
+    basis_func_cls, batch_size
 ):
     """
     Ensure that the random fourier feature map can be used to approximate the covariance matrix
     between the inducing point vectors of a sparse GP, with the condition that the number of latent
     GP models is greater than one.
     """
-    n_features = 10000
+    n_components = 10000
 
     kernel = gpflow.kernels.SquaredExponential()
-    fourier_features = random_feature_cls(kernel, n_features, dtype=tf.float64)
+    fourier_features = basis_func_cls(kernel, n_components, dtype=tf.float64)
 
     x_new = tf.ones(shape=(2 * batch_size + 1, 1), dtype=tf.float64)
 
@@ -118,10 +118,10 @@ def test_fourier_feature_layer_compute_covariance_of_inducing_variables(
     np.testing.assert_allclose(approx_kernel_matrix, actual_kernel_matrix, atol=5e-2)
 
 
-def test_fourier_features_shapes(random_feature_cls, n_components, n_dims, batch_size):
+def test_fourier_features_shapes(basis_func_cls, n_components, n_dims, batch_size):
     input_shape = (batch_size, n_dims)
     kernel = gpflow.kernels.SquaredExponential()
-    feature_functions = random_feature_cls(kernel, n_components, dtype=tf.float64)
+    feature_functions = basis_func_cls(kernel, n_components, dtype=tf.float64)
     output_shape = feature_functions.compute_output_shape(input_shape)
     features = feature_functions(tf.ones(shape=input_shape))
     np.testing.assert_equal(features.shape, output_shape)

From 4b8842e54c21695f95a9e4f87154eb16e6de8a2e Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Tue, 19 Oct 2021 15:57:39 +0100
Subject: [PATCH 13/23] formatting

---
 .../basis_functions/fourier_features/test_quadrature.py     | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
index b7898d6f..d9a1d98b 100644
--- a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
+++ b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
@@ -62,7 +62,9 @@ def test_throw_for_unsupported_kernel():
     assert "Unsupported Kernel" in str(excinfo.value)
 
 
-def test_quadrature_fourier_features_can_approximate_kernel_multidim(kernel_cls, variance, lengthscale, n_dims):
+def test_quadrature_fourier_features_can_approximate_kernel_multidim(
+    kernel_cls, variance, lengthscale, n_dims
+):
 
     n_components = 60
 
@@ -92,7 +94,7 @@ def test_fourier_features_shapes(n_components, n_dims, batch_size):
     feature_functions = QuadratureFourierFeatures(kernel, n_components, dtype=tf.float64)
     output_shape = feature_functions.compute_output_shape(input_shape)
     output_dim = output_shape[-1]
-    assert output_dim == 2*n_components**n_dims
+    assert output_dim == 2 * n_components ** n_dims
     features = feature_functions(tf.ones(shape=input_shape))
     np.testing.assert_equal(features.shape, output_shape)
 

From bc864a7bb3edab2a489fc64d89399b7d4f44ca7e Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Wed, 20 Oct 2021 15:26:56 +0100
Subject: [PATCH 14/23] fixed tests

---
 .../layers/basis_functions/fourier_features/quadrature.py | 7 +++----
 .../basis_functions/fourier_features/test_quadrature.py   | 8 ++++----
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/quadrature.py b/gpflux/layers/basis_functions/fourier_features/quadrature.py
index ddae0e8d..e46ed7f7 100644
--- a/gpflux/layers/basis_functions/fourier_features/quadrature.py
+++ b/gpflux/layers/basis_functions/fourier_features/quadrature.py
@@ -50,9 +50,8 @@ def build(self, input_shape: ShapeType) -> None:
 
         # Quadrature node points
         self.abscissa = tf.Variable(initial_value=abscissa_value, trainable=False)  # (M^D, D)
-        # Gauss-Hermite weights (not to be confused with random Fourier feature
-        # weights or NN weights)
-        self.omegas = tf.Variable(initial_value=omegas_value, trainable=False)  # (M^D,)
+        # Gauss-Hermite weights
+        self.factors = tf.Variable(initial_value=omegas_value, trainable=False)  # (M^D,)
         super(QuadratureFourierFeatures, self).build(input_shape)
 
     def _compute_output_dim(self, input_shape: ShapeType) -> int:
@@ -60,7 +59,7 @@ def _compute_output_dim(self, input_shape: ShapeType) -> int:
         return 2 * self.n_components ** input_dim
 
     def _compute_constant(self) -> tf.Tensor:
-        return tf.tile(tf.sqrt(self.kernel.variance * self.omegas), multiples=[2])
+        return tf.tile(tf.sqrt(self.kernel.variance * self.factors), multiples=[2])
 
     def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
         return _bases_concat(inputs, self.abscissa)
diff --git a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
index d9a1d98b..aa9736ee 100644
--- a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
+++ b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
@@ -40,7 +40,7 @@ def _variance_fixture(request):
     return request.param
 
 
-@pytest.fixture(name="lengthscale", params=[0.1, 1.0, 5.0])
+@pytest.fixture(name="lengthscale", params=[0.75, 1.0, 5.0])
 def _lengthscale_fixture(request):
     return request.param
 
@@ -66,12 +66,12 @@ def test_quadrature_fourier_features_can_approximate_kernel_multidim(
     kernel_cls, variance, lengthscale, n_dims
 ):
 
-    n_components = 60
+    n_components = 128
 
     x_rows = 20
     y_rows = 30
     # ARD
-    lengthscales = np.random.rand(n_dims) * lengthscale
+    lengthscales = np.random.uniform(low=0.75, size=n_dims) * lengthscale
 
     kernel = kernel_cls(variance=variance, lengthscales=lengthscales)
     fourier_features = QuadratureFourierFeatures(kernel, n_components, dtype=tf.float64)
@@ -85,7 +85,7 @@ def test_quadrature_fourier_features_can_approximate_kernel_multidim(
 
     actual_kernel_matrix = kernel.K(x, y)
 
-    np.testing.assert_allclose(approx_kernel_matrix, actual_kernel_matrix, atol=5e-2)
+    np.testing.assert_allclose(approx_kernel_matrix, actual_kernel_matrix)
 
 
 def test_fourier_features_shapes(n_components, n_dims, batch_size):

From 4a7d70fd359ba38d0562bd38d1440a3f9eda02b0 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Wed, 20 Oct 2021 17:47:08 +0100
Subject: [PATCH 15/23] tiny fix

---
 .../notebooks/efficient_posterior_sampling.py | 26 +++++++++----------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/docs/notebooks/efficient_posterior_sampling.py b/docs/notebooks/efficient_posterior_sampling.py
index f4ddd00e..cb9a561b 100644
--- a/docs/notebooks/efficient_posterior_sampling.py
+++ b/docs/notebooks/efficient_posterior_sampling.py
@@ -98,7 +98,7 @@
 num_input_dimensions = [2, 4, 8]  # number of input dimensions
 train_sample_exponents = [2, 4, 6, 8, 10]  # num_train_samples = 2 ** train_sample_exponents
 num_train_samples = [2 ** train_sample_exponent for train_sample_exponent in train_sample_exponents]
-num_components = [
+num_features = [
     1024,
     4096,
     16384,
@@ -233,7 +233,7 @@ def log10_Wasserstein_distance(
 """
 
 # %%
-def conduct_experiment(num_input_dimensions, num_train_samples, num_components):
+def conduct_experiment(num_input_dimensions, num_train_samples, num_features):
     """
     Compute the log10 Wassertein distance between the weight space approximated GP and the exact GP,
     and between the hybrid-rule approximated GP and the exact GP.
@@ -247,7 +247,7 @@ def conduct_experiment(num_input_dimensions, num_train_samples, num_components):
     lengthscale = (
         num_input_dimensions / 100.0
     ) ** 0.5  # adjust kernel lengthscale to the number of input dims
-    num_features = num_train_samples + num_components
+    num_features = num_train_samples + num_features
 
     # exact kernel
     exact_kernel = kernel_class(lengthscales=lengthscale)
@@ -255,7 +255,7 @@ def conduct_experiment(num_input_dimensions, num_train_samples, num_components):
     # weight space approximated kernel
     feature_functions = RandomFourierFeaturesCosine(
         kernel=kernel_class(lengthscales=lengthscale),
-        n_components=num_components,
+        n_components=num_features,
         dtype=default_float(),
     )
     feature_coefficients = np.ones((num_features, 1), dtype=default_float())
@@ -327,7 +327,7 @@ def conduct_experiment(num_input_dimensions, num_train_samples, num_components):
 """
 
 # %%
-def conduct_experiment_for_multiple_runs(num_input_dimensions, num_train_samples, num_components):
+def conduct_experiment_for_multiple_runs(num_input_dimensions, num_train_samples, num_features):
     """
     Conduct the experiment as specified above `num_experiment_runs` times and identify the quartiles for
     the log10 Wassertein distance between the weight space approximated GP and the exact GP,
@@ -347,7 +347,7 @@ def conduct_experiment_for_multiple_runs(num_input_dimensions, num_train_samples
         log10_ws_dist_weight, log10_ws_dist_hybrid = conduct_experiment(
             num_input_dimensions=num_input_dimensions,
             num_train_samples=num_train_samples,
-            num_components=num_components,
+            num_features=num_features,
         )
         list_of_log10_ws_dist_weight.append(log10_ws_dist_weight)
         list_of_log10_ws_dist_hybrid.append(log10_ws_dist_hybrid)
@@ -363,7 +363,7 @@ def conduct_experiment_for_multiple_runs(num_input_dimensions, num_train_samples
 """
 
 # %%
-def conduct_experiment_for_different_train_data_sizes(num_input_dimensions, num_components):
+def conduct_experiment_for_different_train_data_sizes(num_input_dimensions, num_features):
     """
     Conduct the experiment as specified above for different training dataset sizes and store the results in lists.
 
@@ -383,13 +383,13 @@ def conduct_experiment_for_different_train_data_sizes(num_input_dimensions, num_
         ) = conduct_experiment_for_multiple_runs(
             num_input_dimensions=num_input_dimensions,
             num_train_samples=nts,
-            num_components=num_components,
+            num_features=num_features,
         )
         print(
             "Completed for num input dims = "
             + str(num_input_dimensions)
             + " and feature param = "
-            + str(num_components)
+            + str(num_features)
             + " and num train samples = "
             + str(nts)
         )
@@ -420,9 +420,9 @@ def conduct_experiment_for_different_num_features(num_input_dimensions):
         []
     )  # for the analytic solution using the weight space approximated kernel
     list_of_hybrid_results = []  # for the hybrid-rule approximation
-    for nf in num_components:
+    for nf in num_features:
         weight_results, hybrid_results = conduct_experiment_for_different_train_data_sizes(
-            num_input_dimensions=num_input_dimensions, num_components=nf
+            num_input_dimensions=num_input_dimensions, num_features=nf
         )
         print()
         list_of_weight_results.append(weight_results)
@@ -454,7 +454,7 @@ def conduct_experiment_for_different_num_features(num_input_dimensions):
 
     # plot the results for the analytic solution using the weight space approximated kernel
     colors = ["bisque", "orange", "peru"]
-    assert len(colors) == len(num_components), "Number of colors must equal the number of features!"
+    assert len(colors) == len(num_features), "Number of colors must equal the number of features!"
     for j in range(len(weight_results)):
         weight_result = weight_results[j]
         axs[i].fill_between(
@@ -465,7 +465,7 @@ def conduct_experiment_for_different_num_features(num_input_dimensions):
 
     # plot the results for the hybrid-rule approximation
     colors = ["lightblue", "blue", "darkblue"]
-    assert len(colors) == len(num_components), "Number of colors must equal the number of features!"
+    assert len(colors) == len(num_features), "Number of colors must equal the number of features!"
     for j in range(len(hybrid_results)):
         hybrid_result = hybrid_results[j]
         axs[i].fill_between(

From cae35883d53ee8917861fc0f067bfa650e4f719b Mon Sep 17 00:00:00 2001
From: Louis Tiao <1112238+ltiao@users.noreply.github.com>
Date: Wed, 20 Oct 2021 18:17:01 +0100
Subject: [PATCH 16/23] Update
 gpflux/layers/basis_functions/fourier_features/base.py

Co-authored-by: Vincent Dutordoir <dutordoirv@gmail.com>
---
 .../basis_functions/fourier_features/base.py    | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/gpflux/layers/basis_functions/fourier_features/base.py b/gpflux/layers/basis_functions/fourier_features/base.py
index dca96974..084aa66c 100644
--- a/gpflux/layers/basis_functions/fourier_features/base.py
+++ b/gpflux/layers/basis_functions/fourier_features/base.py
@@ -1,3 +1,20 @@
+#
+# Copyright (c) 2021 The GPflux Contributors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+""" Shared functionality for stationary kernel basis functions. """
+
 from abc import ABC, abstractmethod
 from typing import Mapping
 

From a99a23766d4105c0b7260bfdc6a62aca9da67dd9 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Thu, 21 Oct 2021 14:12:24 +0100
Subject: [PATCH 17/23] added final test

---
 .../fourier_features/test_quadrature.py       | 40 +++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
index aa9736ee..4e3d5908 100644
--- a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
+++ b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
@@ -20,6 +20,8 @@
 from tensorflow.python.keras.utils.kernelized_utils import inner_product
 
 import gpflow
+from gpflow.quadrature.gauss_hermite import NDiagGHQuadrature
+from gpflow.utilities.ops import difference_matrix
 
 from gpflux.layers.basis_functions.fourier_features import QuadratureFourierFeatures
 from gpflux.layers.basis_functions.fourier_features.utils import QFF_SUPPORTED_KERNELS
@@ -88,6 +90,44 @@ def test_quadrature_fourier_features_can_approximate_kernel_multidim(
     np.testing.assert_allclose(approx_kernel_matrix, actual_kernel_matrix)
 
 
+def test_feature_map_decomposition(kernel_cls, variance, lengthscale, n_dims, n_components):
+    """
+    Verify that the inner product of the feature map yields exactly the same
+    result as the direct Gauss-Hermite quadrature scheme.
+    """
+    x_rows = 20
+    y_rows = 30
+    # ARD
+    lengthscales = np.random.rand(n_dims) * lengthscale
+
+    kernel = kernel_cls(variance=variance, lengthscales=lengthscales)
+    fourier_features = QuadratureFourierFeatures(kernel, n_components, dtype=tf.float64)
+
+    x = tf.random.uniform((x_rows, n_dims), dtype=tf.float64)
+    y = tf.random.uniform((y_rows, n_dims), dtype=tf.float64)
+
+    u = fourier_features(x)
+    v = fourier_features(y)
+    K_decomp = inner_product(u, v)
+
+    x_scaled = tf.truediv(x, lengthscales)
+    y_scaled = tf.truediv(y, lengthscales)
+    r = difference_matrix(x_scaled, y_scaled)
+
+    def eigen_func(w):
+        return variance * tf.cos(tf.matmul(r, w, transpose_b=True))
+
+    quadrature = NDiagGHQuadrature(dim=n_dims, n_gh=n_components)
+    K_quadrature = quadrature(
+        eigen_func,
+        mean=tf.zeros((1, 1, n_dims), dtype=tf.float64),
+        var=tf.ones((1, 1, n_dims), dtype=tf.float64),
+    )
+    K_quadrature = tf.squeeze(K_quadrature, axis=-1)
+
+    np.testing.assert_allclose(K_decomp, K_quadrature, atol=1e-15)
+
+
 def test_fourier_features_shapes(n_components, n_dims, batch_size):
     input_shape = (batch_size, n_dims)
     kernel = gpflow.kernels.SquaredExponential()

From b395585f3551359aad522f50b2fe3b997f95d207 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Thu, 21 Oct 2021 16:15:17 +0100
Subject: [PATCH 18/23] added warning

---
 gpflux/layers/basis_functions/fourier_features/quadrature.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/gpflux/layers/basis_functions/fourier_features/quadrature.py b/gpflux/layers/basis_functions/fourier_features/quadrature.py
index e46ed7f7..c44123db 100644
--- a/gpflux/layers/basis_functions/fourier_features/quadrature.py
+++ b/gpflux/layers/basis_functions/fourier_features/quadrature.py
@@ -20,6 +20,8 @@
 import tensorflow as tf
 
 import gpflow
+import warnings
+
 from gpflow.base import TensorType
 from gpflow.quadrature.gauss_hermite import ndgh_points_and_weights
 
@@ -34,6 +36,9 @@
 class QuadratureFourierFeatures(FourierFeaturesBase):
     def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
         assert isinstance(kernel, QFF_SUPPORTED_KERNELS), "Unsupported Kernel"
+        if tf.reduce_any(tf.less(kernel.lengthscales, 1e-1)):
+            warnings.warn("Quadrature Fourier feature approximation of kernels "
+                          "with small lengthscale lead to unexpected behaviors!")
         super(QuadratureFourierFeatures, self).__init__(kernel, n_components, **kwargs)
 
     def build(self, input_shape: ShapeType) -> None:

From 353c3d0bcbc04088ef14b2432be20ac7a539369e Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Thu, 21 Oct 2021 16:38:10 +0100
Subject: [PATCH 19/23] formatting

---
 .../basis_functions/fourier_features/quadrature.py       | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/quadrature.py b/gpflux/layers/basis_functions/fourier_features/quadrature.py
index c44123db..9a21d511 100644
--- a/gpflux/layers/basis_functions/fourier_features/quadrature.py
+++ b/gpflux/layers/basis_functions/fourier_features/quadrature.py
@@ -15,13 +15,12 @@
 #
 """ A kernel's features and coefficients using quadrature Fourier features (QFF). """
 
+import warnings
 from typing import Mapping
 
 import tensorflow as tf
 
 import gpflow
-import warnings
-
 from gpflow.base import TensorType
 from gpflow.quadrature.gauss_hermite import ndgh_points_and_weights
 
@@ -37,8 +36,10 @@ class QuadratureFourierFeatures(FourierFeaturesBase):
     def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
         assert isinstance(kernel, QFF_SUPPORTED_KERNELS), "Unsupported Kernel"
         if tf.reduce_any(tf.less(kernel.lengthscales, 1e-1)):
-            warnings.warn("Quadrature Fourier feature approximation of kernels "
-                          "with small lengthscale lead to unexpected behaviors!")
+            warnings.warn(
+                "Quadrature Fourier feature approximation of kernels "
+                "with small lengthscale lead to unexpected behaviors!"
+            )
         super(QuadratureFourierFeatures, self).__init__(kernel, n_components, **kwargs)
 
     def build(self, input_shape: ShapeType) -> None:

From ad6d37027271b8b843ea5e4837ed0944910d932f Mon Sep 17 00:00:00 2001
From: Louis Tiao <1112238+ltiao@users.noreply.github.com>
Date: Fri, 22 Oct 2021 11:41:56 +0100
Subject: [PATCH 20/23] Update
 gpflux/layers/basis_functions/fourier_features/utils.py

Co-authored-by: Vincent Dutordoir <dutordoirv@gmail.com>
---
 gpflux/layers/basis_functions/fourier_features/utils.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/gpflux/layers/basis_functions/fourier_features/utils.py b/gpflux/layers/basis_functions/fourier_features/utils.py
index 950f5847..8c4ff4e1 100644
--- a/gpflux/layers/basis_functions/fourier_features/utils.py
+++ b/gpflux/layers/basis_functions/fourier_features/utils.py
@@ -25,6 +25,12 @@
 
 from gpflux.types import ShapeType
 
+""" 
+Kernels supported by :class:`QuadratureFourierFeatures`. 
+
+Currently we only support the :class:`gpflow.kernels.SquaredExponential` kernel.
+For Matern kernels use :class:`RandomFourierFeatures`.
+"""
 QFF_SUPPORTED_KERNELS: Tuple[Type[gpflow.kernels.Stationary], ...] = (
     gpflow.kernels.SquaredExponential,
 )

From 43556000ab02cf184fcbd27772166240fef6adb2 Mon Sep 17 00:00:00 2001
From: Louis Tiao <1112238+ltiao@users.noreply.github.com>
Date: Fri, 22 Oct 2021 11:42:50 +0100
Subject: [PATCH 21/23] Update
 tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py

Co-authored-by: Vincent Dutordoir <dutordoirv@gmail.com>
---
 .../basis_functions/fourier_features/test_quadrature.py       | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
index 4e3d5908..b1141fba 100644
--- a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
+++ b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
@@ -68,6 +68,10 @@ def test_quadrature_fourier_features_can_approximate_kernel_multidim(
     kernel_cls, variance, lengthscale, n_dims
 ):
 
+    """
+    Compare finite feature approximation to analytical kernel expression.
+    Approximation only holds for large enough lengthscales, as explained in TODO: notebook.
+    """
     n_components = 128
 
     x_rows = 20

From cd8813b0889e387bac92b92002cf75c923022975 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Fri, 22 Oct 2021 14:12:49 +0100
Subject: [PATCH 22/23] added docstrings

---
 .../basis_functions/fourier_features/base.py     |  8 ++++----
 .../fourier_features/quadrature.py               | 16 +++++++++++++---
 .../basis_functions/fourier_features/random.py   | 14 ++++++++++++++
 .../fourier_features/test_quadrature.py          |  2 ++
 4 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/base.py b/gpflux/layers/basis_functions/fourier_features/base.py
index 084aa66c..1aef1bb4 100644
--- a/gpflux/layers/basis_functions/fourier_features/base.py
+++ b/gpflux/layers/basis_functions/fourier_features/base.py
@@ -29,13 +29,13 @@
 class FourierFeaturesBase(ABC, tf.keras.layers.Layer):
     def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
         """
-        :param kernel: kernel to approximate using a set of random features.
-        :param output_dim: total number of basis functions used to approximate
-            the kernel.
+        :param kernel: kernel to approximate using a set of Fourier bases.
+        :param n_components: number of components (e.g. Monte Carlo samples, 
+            quadrature nodes, etc.) used to numerically approximate the kernel.
         """
         super(FourierFeaturesBase, self).__init__(**kwargs)
         self.kernel = kernel
-        self.n_components = n_components  # M: number of Monte Carlo samples
+        self.n_components = n_components
         if kwargs.get("input_dim", None):
             self._input_dim = kwargs["input_dim"]
             self.build(tf.TensorShape([self._input_dim]))
diff --git a/gpflux/layers/basis_functions/fourier_features/quadrature.py b/gpflux/layers/basis_functions/fourier_features/quadrature.py
index 9a21d511..498beff3 100644
--- a/gpflux/layers/basis_functions/fourier_features/quadrature.py
+++ b/gpflux/layers/basis_functions/fourier_features/quadrature.py
@@ -64,8 +64,18 @@ def _compute_output_dim(self, input_shape: ShapeType) -> int:
         input_dim = input_shape[-1]
         return 2 * self.n_components ** input_dim
 
-    def _compute_constant(self) -> tf.Tensor:
-        return tf.tile(tf.sqrt(self.kernel.variance * self.factors), multiples=[2])
-
     def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
+        """
+        Compute basis functions.
+
+        :return: A tensor with the shape ``[N, 2M^D]``.
+        """
         return _bases_concat(inputs, self.abscissa)
+
+    def _compute_constant(self) -> tf.Tensor:
+        """
+        Compute normalizing constant for basis functions.
+
+        :return: A tensor with the shape ``[2M^D,]``
+        """
+        return tf.tile(tf.sqrt(self.kernel.variance * self.factors), multiples=[2])
diff --git a/gpflux/layers/basis_functions/fourier_features/random.py b/gpflux/layers/basis_functions/fourier_features/random.py
index 076dc374..e6b14c70 100644
--- a/gpflux/layers/basis_functions/fourier_features/random.py
+++ b/gpflux/layers/basis_functions/fourier_features/random.py
@@ -113,11 +113,18 @@ def _compute_output_dim(self, input_shape: ShapeType) -> int:
         return 2 * self.n_components
 
     def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
+        """
+        Compute basis functions.
+
+        :return: A tensor with the shape ``[N, 2M]``.
+        """
         return _bases_concat(inputs, self.W)
 
     def _compute_constant(self) -> tf.Tensor:
         """
         Compute normalizing constant for basis functions.
+
+        :return: A tensor with the shape ``[]`` (i.e. a scalar).
         """
         return self.rff_constant(self.kernel.variance, output_dim=2 * self.n_components)
 
@@ -178,10 +185,17 @@ def _compute_output_dim(self, input_shape: ShapeType) -> int:
         return self.n_components
 
     def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
+        """
+        Compute basis functions.
+
+        :return: A tensor with the shape ``[N, M]``.
+        """
         return _bases_cosine(inputs, self.W, self.b)
 
     def _compute_constant(self) -> tf.Tensor:
         """
         Compute normalizing constant for basis functions.
+
+        :return: A tensor with the shape ``[]`` (i.e. a scalar).
         """
         return self.rff_constant(self.kernel.variance, output_dim=self.n_components)
diff --git a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
index b1141fba..c27ce9a8 100644
--- a/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
+++ b/tests/gpflux/layers/basis_functions/fourier_features/test_quadrature.py
@@ -77,6 +77,8 @@ def test_quadrature_fourier_features_can_approximate_kernel_multidim(
     x_rows = 20
     y_rows = 30
     # ARD
+
+    # small lengthscales can lead to large errors in this approximation
     lengthscales = np.random.uniform(low=0.75, size=n_dims) * lengthscale
 
     kernel = kernel_cls(variance=variance, lengthscales=lengthscales)

From 09d27ac9fb0e3847b8d0be40c6be2bc851e2d937 Mon Sep 17 00:00:00 2001
From: Louis Tiao <louistiao@gmail.com>
Date: Fri, 22 Oct 2021 14:13:34 +0100
Subject: [PATCH 23/23] formattingh

---
 gpflux/layers/basis_functions/fourier_features/base.py  | 2 +-
 gpflux/layers/basis_functions/fourier_features/utils.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/gpflux/layers/basis_functions/fourier_features/base.py b/gpflux/layers/basis_functions/fourier_features/base.py
index 1aef1bb4..80461c80 100644
--- a/gpflux/layers/basis_functions/fourier_features/base.py
+++ b/gpflux/layers/basis_functions/fourier_features/base.py
@@ -30,7 +30,7 @@ class FourierFeaturesBase(ABC, tf.keras.layers.Layer):
     def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
         """
         :param kernel: kernel to approximate using a set of Fourier bases.
-        :param n_components: number of components (e.g. Monte Carlo samples, 
+        :param n_components: number of components (e.g. Monte Carlo samples,
             quadrature nodes, etc.) used to numerically approximate the kernel.
         """
         super(FourierFeaturesBase, self).__init__(**kwargs)
diff --git a/gpflux/layers/basis_functions/fourier_features/utils.py b/gpflux/layers/basis_functions/fourier_features/utils.py
index 8c4ff4e1..6959a96c 100644
--- a/gpflux/layers/basis_functions/fourier_features/utils.py
+++ b/gpflux/layers/basis_functions/fourier_features/utils.py
@@ -25,8 +25,8 @@
 
 from gpflux.types import ShapeType
 
-""" 
-Kernels supported by :class:`QuadratureFourierFeatures`. 
+"""
+Kernels supported by :class:`QuadratureFourierFeatures`.
 
 Currently we only support the :class:`gpflow.kernels.SquaredExponential` kernel.
 For Matern kernels use :class:`RandomFourierFeatures`.