docs: def swap_dimensions(method):
      """
      Converts between GPflow indexing and tensorflow indexing
      `method` is a function that broadcasts over the first dimension (i.e. like all tensorflow matrix ops):
          `method` inputs DN1, DNN
          `method` outputs DN1, DNN
      :return: Function that broadcasts over the final dimension (i.e. compatible with GPflow):
          inputs: ND, DNN
          outputs: ND, DNN
      """
      @functools.wraps(method)
      def wrapper(a_nd, b_dnn, swap=True):
          if swap:
              if a_nd.get_shape().ndims != 2:  # pragma: no cover
                  raise ValueError("The `a_nd` input must have 2 dimensions.")
              a_dn1 = tf.transpose(a_nd)[:, :, None]
              A_dn1, B_dnn = method(a_dn1, b_dnn)
              A_nd = tf.transpose(A_dn1[:, :, 0])
              return A_nd, B_dnn
          else:
              return method(a_nd, b_dnn)
      return wrapper