Keras : Simple 1 variable training to get the mean - tensorflow

Question

I wrote a very basic model to train a single variable to approximate the mean value of a vector. But for some reason, it's not working properly.
I used this page describing a linear fit (2 variables):
https://www.tensorflow.org/guide/basic_training_loops
My code is as follow:
import tensorflow as tf
import numpy as np
class MyModel(tf.keras.Model):
def __init__(self, **kwargs):
super().__init__()
self.b = tf.Variable(1.0, trainable=True)
def call(self, x):
return x - self.b
model = MyModel()
model.compile(optimizer=tf.optimizers.Adam(learning_rate=1e-3), loss='mae')
X = np.random.random((10000,1))
Y = np.zeros(X.shape)
model.fit(X, Y, batch_size=10, epochs=10)
B should be optimized so that sum(abs(X - B)) is as close to 0 as possible (= the mean). However when I fit the model it's not training at all and always reaches to the solution B=0 (the real mean is around 0.5).
What do I do wrong?

Answer 1

This code is working fine. Please check below execution and it's output:
import tensorflow as tf
import numpy as np
class MyModel(tf.keras.Model):
def __init__(self, **kwargs):
super().__init__()
self.b = tf.Variable(1.0, trainable=True)
def call(self, x):
return x - self.b
model = MyModel()
model.compile(optimizer=tf.optimizers.Adam(learning_rate=1e-3), loss='mae')
X = np.random.random((10000,1))
Y = np.zeros(X.shape)
model.fit(X, Y, batch_size=10, epochs=10)
Output:
Epoch 1/10
1000/1000 [==============================] - 2s 1ms/step - loss: 0.2991
Epoch 2/10
1000/1000 [==============================] - 1s 1ms/step - loss: 0.2476
Epoch 3/10
1000/1000 [==============================] - 1s 1ms/step - loss: 0.2476
Epoch 4/10
1000/1000 [==============================] - 1s 1ms/step - loss: 0.2476
Epoch 5/10
1000/1000 [==============================] - 1s 1ms/step - loss: 0.2476
Epoch 6/10
1000/1000 [==============================] - 1s 1ms/step - loss: 0.2476
Epoch 7/10
1000/1000 [==============================] - 2s 2ms/step - loss: 0.2476
Epoch 8/10
1000/1000 [==============================] - 2s 2ms/step - loss: 0.2476
Epoch 9/10
1000/1000 [==============================] - 2s 2ms/step - loss: 0.2476
Epoch 10/10
1000/1000 [==============================] - 2s 2ms/step - loss: 0.2477