visualize 10x10 grid of each digit using MNIST samples

I'm trying to plot 10x10 grid samples from the MNIST dataset. 10 of each digit. Here's the code:

Import libraries:

import sklearn
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
from sklearn.pipeline import Pipeline
from sklearn.datasets import fetch_openml

Load digit data:

X, Y = fetch_openml(name='mnist_784', return_X_y=True, cache=False)

Plotting the grid:

def P1(num_examples=10):
plt.rc('image', cmap='Greys')
plt.figure(figsize=(num_examples,len(np.unique(Y))), dpi=X.shape[1])
# For each digit (from 0 to 9)
for i in np.nditer(np.unique(Y)):
    # Create a ndarray with the features of "num_examples" examples of digit "i"
    features = X[Y == i][:num_examples]
    # For each of the "num_examples" examples
    for j in range(num_examples):
        # Create subplot (from 1 to "num_digits"*"num_examples" of each digit)
        plt.subplot(len(np.unique(Y)), num_examples, i * num_examples + j + 1)
        plt.subplots_adjust(wspace=0, hspace=0)
        # Hide tickmarks and scale
        ax = plt.gca()
        # ax.set_axis_off() # Also hide axes (frame) 
        ax.axes.get_xaxis().set_visible(False)
        ax.axes.get_yaxis().set_visible(False)
        # Plot the corresponding digit (reshaped to square matrix/image)
        dim = int(np.sqrt(X.shape[1]))
        digit = features[j].reshape((dim,dim))            
        plt.imshow(digit)
P1(10)

However, I got an error message here saying that: "Iterator operand or requested dtype holds references, but the REFS_OK flag was not enabled"

Can anyone help me with this?

Solution

This error is coming from nd.iter most likely, which you don't need - also recommend using subplots and ax instead of MATLAB style plt calls:

digits = np.unique(Y)
M = 10
dim = int(np.sqrt(X.shape[1]))

fig, axs = plt.subplots(len(digits), M, figsize=(20,20))

for i,d in enumerate(digits):
    for j in range(M):
        axs[i,j].imshow(X[Y==d][j].reshape((dim,dim)))
        axs[i,j].axis('off')