# Buggy code
def K(self, X: tf.Tensor, X2: Optional[tf.Tensor] = None) -> tf.Tensor:
    sig_X = self._sigmoids(X)  # N x 1 x Ncp
    sig_X2 = self._sigmoids(X2) if X2 is not None else sig_X

    # `starters` are the sigmoids going from 0 -> 1, whilst `stoppers` go
    # from 1 -> 0, dimensions are N x N x Ncp
    starters = sig_X * tf.transpose(sig_X2, perm=(1, 0, 2))
    stoppers = (1 - sig_X) * tf.transpose((1 - sig_X2), perm=(1, 0, 2))

    # prepend `starters` with ones and append ones to `stoppers` since the
    # first kernel has no start and the last kernel has no end
    N = tf.shape(X)[0]
    ones = tf.ones((N, N, 1), dtype=X.dtype)
    starters = tf.concat([ones, starters], axis=2)
    stoppers = tf.concat([stoppers, ones], axis=2)

    # now combine with the underlying kernels
    kernel_stack = tf.stack([k(X, X2) for k in self.kernels], axis=2)
    return tf.reduce_sum(kernel_stack * starters * stoppers, axis=2)


# Traceback (InvalidArgumentError)
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-25-d1dbc7941bae> in <module>
----> 1 k(X, xx)

~/Code/GPflow/gpflow/kernels/base.py in __call__(self, X, X2, full_cov, presliced)
170
171         else:
--> 172             return self.K(X, X2)
173
174     def __add__(self, other):

~/Code/GPflow/gpflow/kernels/changepoints.py in K(self, X, X2)
83         N = tf.shape(X)[0]
84         ones = tf.ones((N, N, 1), dtype=X.dtype)
---> 85         starters = tf.concat([ones, starters], axis=2)
86         stoppers = tf.concat([stoppers, ones], axis=2)
87

~/anaconda3/envs/gpflux2/lib/python3.7/site-packages/tensorflow_core/python/util/dispatch.py in wrapper(*args, **kwargs)
178     """Call target, and fall back on dispatchers if there is a TypeError."""
179     try:
--> 180       return target(*args, **kwargs)
181     except (TypeError, ValueError):
182       # Note: convert_to_eager_tensor currently raises a ValueError, not a

~/anaconda3/envs/gpflux2/lib/python3.7/site-packages/tensorflow_core/python/ops/array_ops.py in concat(values, axis, name)
1515           dtype=dtypes.int32).get_shape().assert_has_rank(0)
1516       return identity(values[0], name=name)
-> 1517   return gen_array_ops.concat_v2(values=values, axis=axis, name=name)
1518
1519

~/anaconda3/envs/gpflux2/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_array_ops.py in concat_v2(values, axis, name)
1116         pass  # Add nodes to the TensorFlow graph.
1117     except _core._NotOkStatusException as e:
-> 1118       _ops.raise_from_not_ok_status(e, name)
1119   # Add nodes to the TensorFlow graph.
1120   if not isinstance(values, (list, tuple)):

~/anaconda3/envs/gpflux2/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py in raise_from_not_ok_status(e, name)
6604   message = e.message + (" name: " + name if name is not None else "")
6605   # pylint: disable=protected-access
-> 6606   six.raise_from(core._status_to_exception(e.code, message), None)
6607   # pylint: enable=protected-access
6608

~/anaconda3/envs/gpflux2/lib/python3.7/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: ConcatOp : Dimensions of inputs should match: shape[0] = [100,100,1] vs. shape[1] = [100,25,1] [Op:ConcatV2] name: concat

# Fixed code
def K(self, X: tf.Tensor, X2: Optional[tf.Tensor] = None) -> tf.Tensor:
    sig_X = self._sigmoids(X)  # N1 x 1 x Ncp
    sig_X2 = self._sigmoids(X2) if X2 is not None else sig_X  # N2 x 1 x Ncp

    # `starters` are the sigmoids going from 0 -> 1, whilst `stoppers` go
    # from 1 -> 0, dimensions are N1 x N2 x Ncp
    starters = sig_X * tf.transpose(sig_X2, perm=(1, 0, 2))
    stoppers = (1 - sig_X) * tf.transpose((1 - sig_X2), perm=(1, 0, 2))

    # prepend `starters` with ones and append ones to `stoppers` since the
    # first kernel has no start and the last kernel has no end
    N1 = tf.shape(X)[0]
    N2 = tf.shape(X2)[0] if X2 is not None else N1
    ones = tf.ones((N1, N2, 1), dtype=X.dtype)
    starters = tf.concat([ones, starters], axis=2)
    stoppers = tf.concat([stoppers, ones], axis=2)

    # now combine with the underlying kernels
    kernel_stack = tf.stack([k(X, X2) for k in self.kernels], axis=2)
    return tf.reduce_sum(kernel_stack * starters * stoppers, axis=2)