How to use the tensorly.backend.mean function in tensorly

def MSE(y_true, y_pred, axis=None):
    """Returns the mean squared error between the two predictions

    Parameters
    ----------
    y_true : array of shape (n_samples, )
        Ground truth (correct) target values.
    y_pred : array of shape (n_samples, )
        Estimated target values.

    Returns
    -------
    float
    """
    return T.mean((y_true - y_pred) ** 2, axis=axis)

def covariance(y_true, y_pred, axis=None):
    centered_true = T.mean(y_true, axis=axis)
    centered_pred = T.mean(y_pred, axis=axis)

    if axis is not None:
        # TODO: write a function to do this..
        shape = list(T.shape(y_true))
        shape[axis] = 1
        centered_true = T.reshape(centered_true, shape)
        shape = list(T.shape(y_pred))
        shape[axis] = 1
        centered_pred = T.reshape(centered_pred, shape)

    return T.mean((y_true - centered_true)*(y_pred - centered_pred), axis=axis)

def covariance(y_true, y_pred, axis=None):
    centered_true = T.mean(y_true, axis=axis)
    centered_pred = T.mean(y_pred, axis=axis)

    if axis is not None:
        # TODO: write a function to do this..
        shape = list(T.shape(y_true))
        shape[axis] = 1
        centered_true = T.reshape(centered_true, shape)
        shape = list(T.shape(y_pred))
        shape[axis] = 1
        centered_pred = T.reshape(centered_pred, shape)

    return T.mean((y_true - centered_true)*(y_pred - centered_pred), axis=axis)

def covariance(y_true, y_pred, axis=None):
    centered_true = T.mean(y_true, axis=axis)
    centered_pred = T.mean(y_pred, axis=axis)

    if axis is not None:
        # TODO: write a function to do this..
        shape = list(T.shape(y_true))
        shape[axis] = 1
        centered_true = T.reshape(centered_true, shape)
        shape = list(T.shape(y_pred))
        shape[axis] = 1
        centered_pred = T.reshape(centered_pred, shape)

    return T.mean((y_true - centered_true)*(y_pred - centered_pred), axis=axis)