Свой градиент в TensorFlow

Декоратор custom_gradient оборачивает функцию так, что вместе с прямым значением возвращают правило для обратного прохода. Внутри правила доступны вход и множитель dy сверху по цепочке.

Опишем квадрат входа и явно зададим производную, как при обычном дифференцировании, затем проверим её лентой:

import tensorflow as tf

@tf.custom_gradient
def my_square(x):
	def grad(dy):
		return dy * 2.0 * x
	return x * x, grad

x = tf.Variable(3.0)
with tf.GradientTape() as tape:
	y = my_square(x)
g = tape.gradient(y, x)
print(g) # выведет tf.Tensor(6.0, shape=(), dtype=float32)

Правило можно записать и для ряда: операции идут поэлементно, а множитель dy передаётся в том же виде:

import tensorflow as tf

@tf.custom_gradient
def my_square(x):
	def grad(dy):
		return dy * 2.0 * x
	return x * x, grad

x = tf.Variable([1.0, 2.0, 3.0])
with tf.GradientTape() as tape:
	y = my_square(x)
g = tape.gradient(y, x)
print(g) # выведет tf.Tensor([2. 4. 6.], shape=(3,), dtype=float32)

Опишите удвоение входа с правилом, где обратный множитель проходит без изменений, и выведите производную по переменной 7.0.

Задайте функцию, которая в прямом ходе складывает вход с единицей, а в обратном умножает множитель на два. Проверьте производную по переменной 4.0 на ленте.

Создайте правило для поэлементного утроения аргумента и выведите градиент по ряду 2.0, 5.0.