def mean(X):
    """
    Compute the mean for a dataset of size (D,N) where D is the dimension
    and N is the number of data points
    """
    # given a dataset of size (D, N), the mean should be an array of size (D,1)
    # you can use np.mean, but pay close attention to the
    # shape of the mean vector you are returning.
    D, N = X.shape
    # Edit the code to compute a (D,1) array `mean` for the mean of dataset.
    mean = X.mean(axis=1).reshape(D, 1)
    return mean