詳解ディープラーニングで紹介してる線形と非線形のプロットがうまくいきません

とりあえず投下してみます。

import numpy as np
import tensorflow as tf
from sklearn.utils import shuffle
import matplotlib.pyplot as plt
import matplotlib.animation as animation

np.random.seed(0)
tf.set_random_seed(0)

M = 2      # 入力データの次元
K = 3      # クラス数
n = 100    # クラスごとのデータ数
N = n * K  # 全データ数

X1 = np.random.randn(n, M) + np.array([0, 10])
X2 = np.random.randn(n, M) + np.array([5, 5])
X3 = np.random.randn(n, M) + np.array([10, 0])
Y1 = np.array([[1, 0, 0] for i in range(n)])
Y2 = np.array([[0, 1, 0] for i in range(n)])
Y3 = np.array([[0, 0, 1] for i in range(n)])

X = np.concatenate((X1, X2, X3), axis=0)
Y = np.concatenate((Y1, Y2, Y3), axis=0)
print(X.shape)

W = tf.Variable(tf.zeros([M, K]))
b = tf.Variable(tf.zeros([K]))

x = tf.placeholder(tf.float32, shape=[None, M])
t = tf.placeholder(tf.float32, shape=[None, K])
y = tf.nn.softmax(tf.matmul(x, W) + b)

cross_entropy = tf.reduce_mean(-tf.reduce_sum(t * tf.log(y), axis=1))
train_step = tf.train.GradientDescentOptimizer(0.1).minimize(cross_entropy)

correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(t, 1))

init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)

batch_size = 50  # ミニバッチサイズ
n_batches = N // batch_size

for epoch in range(20):

    X_, Y_ = shuffle(X, Y)

    for i in range(n_batches):
        start = i * batch_size
        end = start + batch_size

        sess.run(train_step, feed_dict={
            x: X_[start:end],
            t: Y_[start:end]
        })

X_, Y_ = shuffle(X, Y)

classified = correct_prediction.eval(session=sess, feed_dict={
    x: X_[0:10],
    t: Y_[0:10]
})
prob = y.eval(session=sess, feed_dict={
    x: X_[0:10]
})

W1 = sess.run(W)
B1 = sess.run(b)

def gen_y(W1, B1):
  W11, W21, W12, W22 = list(W1.ravel())
  B1, B2 = list(B1)
  def y(x):
    return ((W11-W21)*x - (B1-B2)) / (W22 - W12)
  return y
y = gen_y(W1[:, :2], B1[:2])

plt.plot(X1[:,0], X1[:, 1], 'o')
plt.plot(X2[:,0], X2[:, 1], '^')
plt.plot(X3[:,0], X3[:, 1], '*')
plt.xlim((X[:, 0].min(), X[:, 0].max()))
plt.ylim((X[:, 1].min(), X[:, 1].max()))
plt.plot([-2, 12], [y(-2), y(12)])
plt.show()