Repository: GPflow/GPflow

Issue:
fully_correlated_conditional_repeat fvar broadcast_to issue
<!-- Lines like this are comments and will be invisible -->

# Bug

Some parameter combinations for `fully_correlated_conditional_repeat` throw errors:
`q_sqrt = None`
`full_cov` and `full_output_cov` not both `False`.
throw an `InvalidArgumentError`.

## To reproduce

**Minimal, reproducible example**

This is a slightly modified version of the unit test for `fully_correlated_conditional_repeat`.

```python
import gpflow
import tensorflow as tf

Kmm = tf.ones((6, 6)) + gpflow.default_jitter() * tf.eye(6)
Kmn = tf.ones((6, 20, 5))
Knn = tf.ones((20, 5, 5))
f = tf.ones((6, 7))
q_sqrt = None
white = True

m, v = gpflow.conditionals.util.fully_correlated_conditional_repeat(
    Kmn, Kmm, Knn, f, full_cov=False, full_output_cov=True, q_sqrt=q_sqrt, white=white,
)
```

**Stack trace, or error message**
```
/usr/local/lib/python3.7/dist-packages/gpflow/conditionals/util.py in fully_correlated_conditional_repeat(Kmn, Kmm, Knn, f, full_cov, full_output_cov, q_sqrt, white)
    460             fvar = fvar[None, ...] + addvar  # [R, N, P]
    461     else:
--> 462         fvar = tf.broadcast_to(fvar[None], tf.shape(fmean))
    463 
    464     shape_constraints.extend(

/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/gen_array_ops.py in broadcast_to(input, shape, name)
    825       return _result
    826     except _core._NotOkStatusException as e:
--> 827       _ops.raise_from_not_ok_status(e, name)
    828     except _core._FallbackException:
    829       pass

/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name)
   6860   message = e.message + (" name: " + name if name is not None else "")
   6861   # pylint: disable=protected-access
-> 6862   six.raise_from(core._status_to_exception(e.code, message), None)
   6863   # pylint: enable=protected-access
   6864 

/usr/local/lib/python3.7/dist-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: Rank of input (4) must be no greater than rank of output shape (3). [Op:BroadcastTo]
```

## Expected behavior

Return the mean and variance.

## System information

* GPflow version: 2.1.4
* GPflow installed from: git commit 405eb97dd30cb43f02501b58b581f2608eb6e43e
* TensorFlow version: 2.3.0
* Python version: 3.7.5
* Operating system: Ubuntu 18.04.5 LTS

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:

Resolving patch:
diff --git a/gpflow/conditionals/util.py b/gpflow/conditionals/util.py
--- a/gpflow/conditionals/util.py
+++ b/gpflow/conditionals/util.py
@@ -459,7 +459,8 @@ def fully_correlated_conditional_repeat(
             addvar = tf.reshape(tf.reduce_sum(tf.square(LTA), axis=1), (R, N, P))  # [R, N, P]
             fvar = fvar[None, ...] + addvar  # [R, N, P]
     else:
-        fvar = tf.broadcast_to(fvar[None], tf.shape(fmean))
+        fvar_shape = tf.concat([[R], tf.shape(fvar)], axis=0)
+        fvar = tf.broadcast_to(fvar[None], fvar_shape)
 
     shape_constraints.extend(
         [(Knn, intended_cov_shape), (fmean, ["R", "N", "P"]), (fvar, ["R"] + intended_cov_shape),]


Test patch:
diff --git a/tests/gpflow/conditionals/test_multioutput.py b/tests/gpflow/conditionals/test_multioutput.py
--- a/tests/gpflow/conditionals/test_multioutput.py
+++ b/tests/gpflow/conditionals/test_multioutput.py
@@ -317,29 +317,116 @@ def test_sample_conditional_mixedkernel():
     )
 
 
+@pytest.fixture(
+    name="fully_correlated_q_sqrt_factory",
+    params=[lambda _, __: None, lambda LM, R: tf.eye(LM, batch_shape=(R,))],
+)
+def _q_sqrt_factory_fixture(request):
+    return request.param
+
+
+@pytest.mark.parametrize("R", [1, 2, 5])
 @pytest.mark.parametrize(
-    "func, R",
+    "whiten",
     [
-        (fully_correlated_conditional_repeat, 5),
-        (fully_correlated_conditional_repeat, 1),
-        (fully_correlated_conditional, 1),
+        True,
+        pytest.param(
+            False,
+            marks=pytest.mark.xfail(
+                reason="fully_correlated_conditional_repeat does not support whiten=False"
+            ),
+        ),
     ],
 )
-def test_fully_correlated_conditional_repeat_shapes(func, R):
+def test_fully_correlated_conditional_repeat_shapes_fc_and_foc(
+    R, fully_correlated_q_sqrt_factory, full_cov, full_output_cov, whiten
+):
+
     L, M, N, P = Data.L, Data.M, Data.N, Data.P
 
     Kmm = tf.ones((L * M, L * M)) + default_jitter() * tf.eye(L * M)
     Kmn = tf.ones((L * M, N, P))
-    Knn = tf.ones((N, P))
+
+    if full_cov and full_output_cov:
+        Knn = tf.ones((N, P, N, P))
+        expected_v_shape = [R, N, P, N, P]
+    elif not full_cov and full_output_cov:
+        Knn = tf.ones((N, P, P))
+        expected_v_shape = [R, N, P, P]
+    elif full_cov and not full_output_cov:
+        Knn = tf.ones((P, N, N))
+        expected_v_shape = [R, P, N, N]
+    else:
+        Knn = tf.ones((N, P))
+        expected_v_shape = [R, N, P]
+
     f = tf.ones((L * M, R))
-    q_sqrt = None
-    white = True
+    q_sqrt = fully_correlated_q_sqrt_factory(L * M, R)
+
+    m, v = fully_correlated_conditional_repeat(
+        Kmn,
+        Kmm,
+        Knn,
+        f,
+        full_cov=full_cov,
+        full_output_cov=full_output_cov,
+        q_sqrt=q_sqrt,
+        white=whiten,
+    )
+
+    assert m.shape.as_list() == [R, N, P]
+    assert v.shape.as_list() == expected_v_shape
+
 
-    m, v = func(
-        Kmn, Kmm, Knn, f, full_cov=False, full_output_cov=False, q_sqrt=q_sqrt, white=white,
+@pytest.mark.parametrize(
+    "whiten",
+    [
+        True,
+        pytest.param(
+            False,
+            marks=pytest.mark.xfail(
+                reason="fully_correlated_conditional does not support whiten=False"
+            ),
+        ),
+    ],
+)
+def test_fully_correlated_conditional_shapes_fc_and_foc(
+    fully_correlated_q_sqrt_factory, full_cov, full_output_cov, whiten
+):
+    L, M, N, P = Data.L, Data.M, Data.N, Data.P
+
+    Kmm = tf.ones((L * M, L * M)) + default_jitter() * tf.eye(L * M)
+    Kmn = tf.ones((L * M, N, P))
+
+    if full_cov and full_output_cov:
+        Knn = tf.ones((N, P, N, P))
+        expected_v_shape = [N, P, N, P]
+    elif not full_cov and full_output_cov:
+        Knn = tf.ones((N, P, P))
+        expected_v_shape = [N, P, P]
+    elif full_cov and not full_output_cov:
+        Knn = tf.ones((P, N, N))
+        expected_v_shape = [P, N, N]
+    else:
+        Knn = tf.ones((N, P))
+        expected_v_shape = [N, P]
+
+    f = tf.ones((L * M, 1))
+    q_sqrt = fully_correlated_q_sqrt_factory(L * M, 1)
+
+    m, v = fully_correlated_conditional(
+        Kmn,
+        Kmm,
+        Knn,
+        f,
+        full_cov=full_cov,
+        full_output_cov=full_output_cov,
+        q_sqrt=q_sqrt,
+        white=whiten,
     )
 
-    assert v.shape.as_list() == m.shape.as_list()
+    assert m.shape.as_list() == [N, P]
+    assert v.shape.as_list() == expected_v_shape
 
 
 # ------------------------------------------
diff --git a/tests/gpflow/conftest.py b/tests/gpflow/conftest.py
new file mode 100644
--- /dev/null
+++ b/tests/gpflow/conftest.py
@@ -0,0 +1,25 @@
+#  Copyright 2021 The GPflow Contributors. All Rights Reserved.
+#
+#  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 pytest
+
+
+@pytest.fixture(name="full_cov", params=[True, False])
+def _full_cov_fixture(request):
+    return request.param
+
+
+@pytest.fixture(name="full_output_cov", params=[True, False])
+def _full_output_cov_fixture(request):
+    return request.param