Small changes (docstring mistake and unit test improvement)

gpflux/layers/basis_functions/fourier_features/base.py

@@ -71,7 +71,7 @@ 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]``, or shape ``[P, N, M]'' in the batched case.
+        :return: A tensor with the shape ``[N, M]``, or shape ``[P, N, M]'' in the multioutput case.
         """
         if self.is_batched:
             X = [tf.divide(inputs, k.lengthscales) for k in self.sub_kernels]
@@ -82,9 +82,9 @@ def call(self, inputs: TensorType) -> tf.Tensor:
         bases = self._compute_bases(X)  # [N, M] or [P, N, M]
         output = const * bases
 
-        # For combination kernels, remove batch dimension and instead concatenate into the
-        # feature dimension.
         if self.is_batched and not self.is_multioutput:
+            # For combination kernels, remove batch dimension and instead concatenate into the
+            # feature dimension.
             output = tf.transpose(output, perm=[1, 2, 0])  # [N, M, P]
             output = tf.reshape(output, [tf.shape(output)[0], -1])  # [N, M*P]
 

gpflux/layers/basis_functions/fourier_features/random/base.py

@@ -200,10 +200,10 @@ class RandomFourierFeatures(RandomFourierFeaturesBase):
     """
 
     def compute_output_dim(self, input_shape: ShapeType) -> int:
-        # For combination kernels, the number of features is multiplied by the number of
-        # sub-kernels.
         dim = 2 * self.n_components
         if self.is_batched and not self.is_multioutput:
+            # For combination kernels, the number of features is multiplied by the number of
+            # sub-kernels.
             dim *= self.batch_size
         return dim
 
@@ -287,10 +287,10 @@ def _bias_init(self, shape: TensorType, dtype: Optional[DType] = None) -> Tensor
         return tf.random.uniform(shape=shape, maxval=2.0 * np.pi, dtype=dtype)
 
     def compute_output_dim(self, input_shape: ShapeType) -> int:
-        # For combination kernels, the number of features is multiplied by the number of
-        # sub-kernels.
         dim = self.n_components
         if self.is_batched and not self.is_multioutput:
+            # For combination kernels, the number of features is multiplied by the number of
+            # sub-kernels.
             dim *= self.batch_size
         return dim
 

tests/gpflux/layers/basis_functions/fourier_features/test_random.py

@@ -71,7 +71,13 @@ def _kernel_cls_fixture(request):
                 gpflow.kernels.Matern32(lengthscales=0.1),
             ]
         ),
-        gpflow.kernels.Sum([gpflow.kernels.SquaredExponential(), gpflow.kernels.Matern52()]),
+        gpflow.kernels.Sum(
+            [
+                gpflow.kernels.SquaredExponential(),
+                gpflow.kernels.Matern32(),
+                gpflow.kernels.Matern52(),
+            ]
+        ),
     ],
 )
 def _multi_kernel_cls_fixture(request):