# Buggy code
def wrapper(old_fun):
    def new_fun(X, **Ys):
        Xs = tf.unstack(X, axis=-1)
        fun_eval = old_fun(*Xs, **Ys)
        if tf.rank(fun_eval) < tf.rank(X):
            fun_eval = tf.expand_dims(fun_eval, axis=-1)
        return fun_eval

    return new_fun


# Traceback (tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError)
Traceback (most recent call last):
File "gpflow_error.py", line 20, in <module>
go()
File "gpflow_error.py", line 16, in go
quad = compute()
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 580, in __call__
result = self._call(*args, **kwds)
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 627, in _call
self._initialize(args, kwds, add_initializers_to=initializers)
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 506, in _initialize
*args, **kwds))
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2446, in _get_concrete_function_internal_garbage_collected
graph_function, _, _ = self._maybe_define_function(args, kwargs)
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2777, in _maybe_define_function
graph_function = self._create_graph_function(args, kwargs)
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2667, in _create_graph_function
capture_by_value=self._capture_by_value),
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 981, in func_graph_from_py_func
func_outputs = python_func(*func_args, **func_kwargs)
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 441, in wrapped_fn
return weak_wrapped_fn().__wrapped__(*args, **kwds)
File "gpflow_error.py", line 11, in compute
quad = quadrature.ndiagquad([lambda *X: tf.exp(X[0])], num_gauss_hermite_points, [mu], [var])
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/gpflow/quadrature/deprecated.py", line 156, in ndiagquad
result = quadrature(funcs, Fmu, Fvar, **Ys)
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/gpflow/quadrature/base.py", line 52, in __call__
return [tf.reduce_sum(f(X, *args, **kwargs) * W, axis=-2) for f in fun]
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/gpflow/quadrature/base.py", line 52, in <listcomp>
return [tf.reduce_sum(f(X, *args, **kwargs) * W, axis=-2) for f in fun]
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/gpflow/quadrature/deprecated.py", line 141, in new_fun
if tf.rank(fun_eval) < tf.rank(X):
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 778, in __bool__
self._disallow_bool_casting()
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 542, in _disallow_bool_casting
"using a `tf.Tensor` as a Python `bool`")
File "/Users/nferguson/gpflow_error/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 527, in _disallow_when_autograph_disabled
" Try decorating it directly with @tf.function.".format(task))
tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph is disabled in this function. Try decorating it directly with @tf.function.

# Fixed code
def wrapper(old_fun):
    def new_fun(X, **Ys):
        Xs = tf.unstack(X, axis=-1)
        fun_eval = old_fun(*Xs, **Ys)
        return tf.cond(
            pred=tf.less(tf.rank(fun_eval), tf.rank(X)),
            true_fn=lambda: fun_eval[..., tf.newaxis],
            false_fn=lambda: fun_eval,
        )

    return new_fun