x1 = tf.random.uniform([150], minval=0, maxval=10)
x2 = tf.random.uniform([150], minval=0, maxval=10)
y = 1.5 + 2.0 * x1 + 5.0 * x2
dataset = tf.stack([y, x1, x2], axis=1)

dataset_shuffled = tf.random.shuffle(dataset, seed=22)
train_data, test_data = dataset_shuffled[:125], dataset_shuffled[125:]
x_train, y_train = train_data[:, 1:], train_data[:, 0]

class LinearRegression(tf.Module):
    @tf.function
    def __call__(self, x):
        if not self.built:
            self.w = tf.Variable(tf.random.uniform([x.shape[-1], 1]))
            self.b = tf.Variable(tf.random.uniform([]))
            self.built = True
        return tf.squeeze(tf.matmul(x, self.w) + self.b, axis=1)

def mse_loss(y_pred, y):
    return tf.reduce_mean(tf.square(y_pred - y))

train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_dataset = train_dataset.shuffle(buffer_size=125).batch(32)

for epoch in range(epochs):
    for x_batch, y_batch in train_dataset:
        with tf.GradientTape() as tape:
            y_pred = lin_reg(x_batch)
            loss = mse_loss(y_pred, y_batch)
        grads = tape.gradient(loss, lin_reg.variables)
        for g, v in zip(grads, lin_reg.variables):
            v.assign_sub(learning_rate * g)