import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import tensorflow as tf

x = tf.constant(3.0)

with tf.GradientTape() as tape:
    tape.watch(x)
    y = x * x + 2.0 * x

dy_dx = tape.gradient(y, x)

weight = tf.Variable(2.0, name="weight")
feature = tf.constant(3.0)
target = tf.constant(15.0)
optimizer = tf.keras.optimizers.SGD(learning_rate=0.05)

with tf.GradientTape() as tape:
    prediction = weight * feature
    loss_before = tf.square(prediction - target)

gradient = tape.gradient(loss_before, weight)

if gradient is None:
    raise RuntimeError("The loss is not connected to weight.")

weight_before = tf.identity(weight)
optimizer.apply_gradients([(gradient, weight)])
loss_after = tf.square(weight * feature - target)

tf.debugging.assert_near(dy_dx, tf.constant(8.0))
tf.debugging.assert_less(loss_after, loss_before)

print(f"dy/dx at x=3.0: {float(dy_dx):.2f}")
print(f"weight gradient: {float(gradient):.2f}")
print(f"weight: {float(weight_before):.2f} -> {float(weight):.2f}")
print(f"loss: {float(loss_before):.2f} -> {float(loss_after):.2f}")