encoded_set = large.encode(x_test)
encoded_test = tf.data.Dataset.from_tensor_slices((encoded_set, y_test))
encoded_sub = encoded_test.take(20000)
data = ([],[],[],[],[],[],[],[],[],[])
shapes = (".", ",", "o", "v", "^", "<", ">", "1", "2", "3")
labels = ("9", "8", "7", "6", "5", "4", "3", "2", "1", "0")


for input, output in iter(encoded_sub):
  if output.numpy() == 9:
    data[9].append(input.numpy())
  elif output.numpy() == 8:
    data[8].append(input.numpy())
  elif output.numpy() == 7:
    data[7].append(input.numpy())
  elif output.numpy() == 6:
    data[6].append(input.numpy())
  elif output.numpy() == 5:
    data[5].append(input.numpy())
  elif output.numpy() == 4:
    data[4].append(input.numpy())
  elif output.numpy() == 3:
    data[3].append(input.numpy())
  elif output.numpy() == 2:
    data[2].append(input.numpy())
  elif output.numpy() == 1:
    data[1].append(input.numpy())
  elif output.numpy() == 0:
    data[0].append(input.numpy())


fig = plt.figure()
ax1 = fig.add_subplot(111)
i = 0
for sets in data:
  x_cords = []
  y_cords = []
  for item in sets:
    x_cords.append(item[0])
    y_cords.append(item[1])


  ax1.scatter(x_cords, y_cords, marker = shapes[i], label = labels[i])
  i+=1
ax1.legend()