Методы оптимизации из TensorFlow своими руками

Требуется выполнить оптимизацию - обучить модель методом nesterov momentum и rmsprop. Сделал, но получаю ошибки размерности. В чем затык?

import pandas as pd
import numpy as np
from sklearn.datasets import load_iris

X,y = load_iris(return_X_y=True)
filter = y != 2              # накладываем фильтр - только 2 признака
X = X[filter]
y = y[filter]

COEFS = np.random.randn(5)   # генерируем случайные коэффициенты

def predict_proba(coefs, x):
  # формула логистической регрессии:
  return 1. / (1. + np.exp( -(x.dot(coefs[:4]) + coefs[-1]) ) ) 

# теперь на основе модели предсказываем класс (но модель пока не обучена: 
# это сделаем чуть позже)
def predict_class(coefs, x):
  probas = predict_proba(coefs, x)
  return (probas > 0.5).astype(np.float)

# явно прописываем функцию потерь на основе ее формулы
def bce_loss(coefs, x, y):
  probas = predict_proba(coefs, x)
  filter_ones = y == 1
  loss = -1. * (np.sum(np.log(probas[filter_ones])) + np.sum(np.log(1. - probas[~filter_ones]))) / len(y)
  return loss

# расчет градиента:
# он зависит от двух сущностей: от модели и функции потерь
def grad(coefs, x, y):
  probas = predict_proba(coefs, x)
  delta = probas - y
  modified_x = x.T * delta
  deltas = np.mean(modified_x, axis=1)
  return deltas, np.mean(delta)

def learn_nesterov(coefs, x, y, num_epochs=20, momentum=0.9, learning_rate=0.0001):
    v_t = [0 for _ in range(len(coefs))]
    losses =[]
    for it in range(num_epochs):
        pr_coefs  = [coefs[i] - momentum*v_t[i] for i in range(len(coefs))]
        gr_coefs = grad(pr_coefs, x, y)
        for i in range(len(coefs)):
            v_t[i]=momentum*v_t[i] + learning_rate*gr_coefs[i]
            coefs[i] = coefs[i]-v_t[i]
        losses.append(bce_loss(x, y, coefs))
    return losses, coefs

learn_nesterov(COEFS, X, y)

Таким образом, выше реализация эвристики на основе метода Нестерова. Вторая эвристика ниже

def learn_RMSprop(coefs, x, y, num_epochs=20, momentum=0.9, learning_rate=0.0001):
    e = 10**(-8)
    S = [0 for _ in range(len(coefs))]
    losses =[]
    for it in range(num_epochs):
        gr_coefs = grad(x, y, coefs)
        gr_coefs_2 = [x**2 for x in gr_coefs]
        for i in range(len(coefs)):
            S[i]=momentum*S[i] + (1-momentum)*gr_coefs_2[i]
            coefs[i] = coefs[i]-learning_rate* (gr_coefs[i]/(math.sqrt(S[i])+e))  
        losses.append(bce_loss(coefs, x, y))
    return losses, coefs

learn_RMSprop(COEFS, X, y)