Generate the Dataset

import tensorflow as tf
import random
import matplotlib.pyplot as plt
%matplotlib inline
def synthetic_data(w, b, num_examples):  
    """Generate y = Xw + b + noise."""
    X = tf.zeros((num_examples, w.shape[0]))
    X += tf.random.normal(shape=X.shape)
    y = tf.matmul(X, tf.reshape(w, (-1, 1))) + b
    y += tf.random.normal(shape=y.shape, stddev=0.01)
    y = tf.reshape(y, (-1, 1))
    return X, y

true_w = tf.constant([2, -3.4])
true_b = 4.2
X, y = synthetic_data(true_w, true_b, 1000)

Reading the Dataset

def load_array(data_arrays, batch_size, is_train=True):  #@save
    """Construct a TensorFlow data iterator."""
    dataset = tf.data.Dataset.from_tensor_slices(data_arrays)
    if is_train:
        dataset = dataset.shuffle(buffer_size=1000)
    dataset = dataset.batch(batch_size)
    return dataset

batch_size = 10
data_iter = load_array((X, y), batch_size)

Defining the Model & Initializing Parameters

initializer = tf.initializers.RandomNormal(stddev=0.01)
net = tf.keras.Sequential()
net.add(tf.keras.layers.Dense(1, kernel_initializer=initializer))

Defining the Loss Function

loss = tf.keras.losses.MeanSquaredError()
loss = tf.keras.losses.Huber()

Defining the Optimization Algorithm

trainer = tf.keras.optimizers.SGD(learning_rate=0.03)

Training

num_epochs = 3
for epoch in range(num_epochs):
    for X_,y_ in data_iter:
        with tf.GradientTape() as tape:
            l = loss(y_,net(X_, training = True))
        grads = tape.gradient(l,net.trainable_variables)
        trainer.apply_gradients(zip(grads,net.trainable_variables))
    l = loss(net(X), y)
    print(f'epoch {epoch +1}, loss {l:f}')
epoch 1, loss 2.276382
epoch 2, loss 0.314810
epoch 3, loss 0.000611
net.get_weights()
[array([[ 1.9885834],
        [-3.379021 ]], dtype=float32),
 array([4.1781187], dtype=float32)]
a = tf.reshape(tf.range(3), (3, 1))
b = tf.reshape(tf.range(2), (1, 2))
a, b
(<tf.Tensor: shape=(3, 1), dtype=int32, numpy=
 array([[0],
        [1],
        [2]], dtype=int32)>,
 <tf.Tensor: shape=(1, 2), dtype=int32, numpy=array([[0, 1]], dtype=int32)>)