Мультиклассовая логистическая регрессия

Я пытаюсь написать многоклассовую логистическую регрессию с помощью numpy и выполнить перекрестное проверочное тестирование с использованием MNIST:

num_classes = 10
num_features = 784
learning_rate = 0.0001
training_steps = 50
batch_size = 256
display_step = 50

from tensorflow.keras.datasets import mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train, x_test = np.array(x_train, np.float32), np.array(x_test, np.float32)

x_train, x_test = x_train.reshape([-1, num_features]), x_test.reshape([-1, num_features])

x_train, x_test = x_train/255., x_test/255.

def logistic_regression(x, b, b0):
  return 1./(1.+np.exp(-np.dot(b,x)-b0))

b = np.random.uniform(-1, 1, num_features*num_classes).reshape((num_classes, num_features))
b0 = np.random.uniform(-1, 1, num_classes)

for step in range(training_steps):
  db = np.zeros((num_classes, num_features), dtype = 'float32')
  db0 = np.zeros(num_classes, dtype = 'float32')

  for x, y in zip (x_train, y_train):
    yy = tf.one_hot(y, depth=num_classes).numpy()
    a = logistic_regression(x, b, b0)
    db += np.matmul(np.expand_dims(yy - a, axis =- 1), np.expand_dims(x, axis = 0))
    db0 += yy - a

  b += learning_rate * db
  b0 += learning_rate * db0

def accuracy(y_pred, y_true):
  correct_prediction = np.equal(np.argmax(y_pred, axis = 1), y_true.astype(np.int64))
  return np.mean(correct_prediction.astype(np.float32))

pred = np.vectorize(logistic_regression, signature='(n),(n),()->()')(x_test, b, b0)

Но приведенный выше код выдает мне следующую ошибку:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-12-a37a0ec6d637> in <module>()
----> 1 pred = np.vectorize(logistic_regression, signature='(n),(n),()->()')(x_test, b, b0)

4 frames
/usr/local/lib/python3.7/dist-packages/numpy/lib/function_base.py in __call__(self, *args, **kwargs)
   2106             vargs.extend([kwargs[_n] for _n in names])
   2107 
-> 2108         return self._vectorize_call(func=func, args=vargs)
   2109 
   2110     def _get_ufunc_and_otypes(self, func, args):

/usr/local/lib/python3.7/dist-packages/numpy/lib/function_base.py in _vectorize_call(self, func, args)
   2180         """Vectorized call to `func` over positional `args`."""
   2181         if self.signature is not None:
-> 2182             res = self._vectorize_call_with_signature(func, args)
   2183         elif not args:
   2184             res = func()

/usr/local/lib/python3.7/dist-packages/numpy/lib/function_base.py in _vectorize_call_with_signature(self, func, args)
   2210 
   2211         broadcast_shape, dim_sizes = _parse_input_dimensions(
-> 2212             args, input_core_dims)
   2213         input_shapes = _calculate_shapes(broadcast_shape, dim_sizes,
   2214                                          input_core_dims)

/usr/local/lib/python3.7/dist-packages/numpy/lib/function_base.py in _parse_input_dimensions(args, input_core_dims)
   1875         dummy_array = np.lib.stride_tricks.as_strided(0, arg.shape[:ndim])
   1876         broadcast_args.append(dummy_array)
-> 1877     broadcast_shape = np.lib.stride_tricks._broadcast_shape(*broadcast_args)
   1878     return broadcast_shape, dim_sizes
   1879 

/usr/local/lib/python3.7/dist-packages/numpy/lib/stride_tricks.py in _broadcast_shape(*args)
    187     # use the old-iterator because np.nditer does not handle size 0 arrays
    188     # consistently
--> 189     b = np.broadcast(*args[:32])
    190     # unfortunately, it cannot handle 32 or more arguments directly
    191     for pos in range(32, len(args), 31):

ValueError: shape mismatch: objects cannot be broadcast to a single shape

Подскажите, пожалуйста, что не так с кодом?