# Source: content/notes/libraries/tensorflow.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import numpy as np
import tensorflow as tf
tf.config.threading.set_intra_op_parallelism_threads(1)
tf.config.threading.set_inter_op_parallelism_threads(1)
tf.keras.utils.set_random_seed(7)
x = tf.Variable(2.)
with tf.GradientTape() as outer:
    with tf.GradientTape() as inner:
        y = x**3
    first = inner.gradient(y, x)
second = outer.gradient(first, x)
assert first.numpy() == 12 and second.numpy() == 12
constant = tf.constant(3.)
with tf.GradientTape() as tape:
    tape.watch(constant)
    square = constant**2
assert tape.gradient(square, constant).numpy() == 6

def update(weight, inputs, targets):
    with tf.GradientTape() as tape:
        prediction = inputs @ weight
        loss = tf.reduce_mean(tf.square(targets-prediction))
    weight.assign_sub(0.05*tape.gradient(loss, weight))
    return loss

inputs = tf.constant([[1., 2.], [3., -1.]])
targets = tf.constant([[1.], [0.]])
a, b = tf.Variable([[0.1], [0.2]]), tf.Variable([[0.1], [0.2]])
eager_loss = update(a, inputs, targets)
graph_loss = tf.function(update)(b, inputs, targets)
np.testing.assert_allclose(a.numpy(), b.numpy(), rtol=1e-6, atol=1e-7)
np.testing.assert_allclose(eager_loss.numpy(), graph_loss.numpy(), rtol=1e-6)
print("watched constants, second derivatives, eager/graph update parity passed")