Linear regression in TF2

linregtf.py

class SimpleLinearRegression:
    def __init__(self, initializer='random'):
        if initializer=='ones':
            self.var = 1.
        elif initializer=='zeros':
            self.var = 0.
        elif initializer=='random':
            selfx.var = tf.random.uniform(shape=[], minval=0., maxval=1.)
            
        self.m = tf.Variable(1., shape=tf.TensorShape(None))
        self.b = tf.Variable(self.var)
        
    def predict(self, x):
        return tf.reduce_sum(self.m * x, 1) + self.b
    
    def mse(self, true, predicted):
        return tf.reduce_mean(tf.square(true-predicted))
    
    def update(self, X, y, learning_rate):
        with tf.GradientTape(persistent=True) as g:
            loss = self.mse(y, self.predict(X))
            
        print("Loss: ", loss)

        dy_dm = g.gradient(loss, self.m)
        dy_db = g.gradient(loss, self.b)
        
        self.m.assign_sub(learning_rate * dy_dm)
        self.b.assign_sub(learning_rate * dy_db)
    
    def train(self, X, y, learning_rate=0.01, epochs=5):
        
        if len(X.shape)==1:
            X=tf.reshape(X,[X.shape[0],1])
        
        self.m.assign([self.var]*X.shape[-1])
        
        for i in range(epochs):
            print("Epoch: ", i)
            
            self.update(X, y, learning_rate)