La regresión lineal simple en Tensorflow produce un coeficiente cercano a cero.

Estoy intentando una regresión lineal simple en Tensorflow con solo una variable independiente. Una trama de mis datos muestra que el coeficiente debería estar cerca de 1, y de hecho, si lo corro usando sklearn.linear_model.LinearRegression, obtengo un resultado sensato de alrededor de 0,90.

Sin embargo, al ejecutarlo en Tensorflow usando este tutorial, produce un coeficiente muy cercano a cero. Pude obtener un resultado racional de Tensorflow usando números aleatorios. He intentado ajustar la tasa de aprendizaje o el número de épocas sin ningún efecto significativo.

El MRE incluye datos reales y debería producir un coeficiente de 0.8975 de sklearn pero 0.00045 de Tensorflow. He considerado que está atrapado en un mínimo local, pero ninguno de los ejemplos que puedo encontrar de tal problema funciona para mi problema.

import numpy as np
import tensorflow as tf
from sklearn import linear_model

learning_rate = 0.1
epochs = 100

x_train = np.array([-0.00055, 0.00509, -0.0046, -0.01687, -0.0047, 0.00348, 
                0.00042, -0.00208, -0.01207, -0.0007, 0.00408, -0.00182, 
                -0.00294, -0.00113, 0.0038, -0.00645, 0.00113, 0.00268, 
                -0.0045, -0.00381, 0.00298, 0, -0.00184, -0.00212, 
                -0.00213, -0.01224, 0.00072, 0, -0.00331, 0.00534, 
                0.00675, -0.00285, -0.00429, 0.00489, -0.00286, 0.00158, 
                0.00129, 0.00472, 0.00555, -0.00467, -0.00231, -0.00231, 
                0.00159, -0.00463, 0.00174, 0, -0.0029, 
                -0.00349, 0.01372, -0.00302])

y_train = np.array([0.00125, 0.00218, -0.00373, -0.00999, -0.00441, 
                0.00412, 0.00158, -0.00094, -0.01513, -0.00064, 0.00416, 
                -0.00191, -0.00607, 0.00161, 0.00289, -0.00416, 
                0.00096, 0.00321, -0.00672, -0.0029, 0.00129, -0.00032, 
                -0.00387, -0.00162, -0.00292, -0.01367, 0.00198, 
                0.00099, -0.00329, 0.00693, 0.00459, -0.00294, -0.00164, 
                0.00328, -0.00425, 0.00131, 0.00131, 0.00524, 0.00358,
                -0.00422, -0.00065, -0.00359, 0.00229, 0, 0.00196, 
                -0.00065, -0.00391, -0.0108, 0.01291, -0.00098])

regr = linear_model.LinearRegression()
regr.fit(x_train.reshape(-1, 1), y_train.reshape(-1, 1))
print('Coefficients: ', regr.coef)

weight = tf.Variable(0.)
bias = tf.Variable(0.)

for e in range(epochs):
    with tf.GradientTape() as tape:
        y_pred = weight*x_train + bias
        loss = tf.reduce_mean(tf.square(y_pred - y_train))
        gradients = tape.gradient(loss, [weight,bias])
        weight.assign_sub(gradients[0]*learning_rate)
        bias.assign_sub(gradients[1]*learning_rate)

print(weight.numpy(), 'weight', bias.numpy(), 'bias')