Как я могу сделать обратную игрушку модель тензор потока с GD

Question

Я пытаюсь написать код тензорного потока, содержащий модель (нейронную сеть с tf.keras.layers) и обучаемую переменную с именем pi. Цель состоит в том, чтобы оптимизировать для pi и для параметров модели через GD.

В приведенном ниже примере с игрушкой потеря = потеря MSE + пи ^ 2. Я ожидаю, что пи сходится к 0.

import tensorflow as tf
import numpy as np

optimizer =    tf.keras.optimizers.Adam
dataset_train = np.array([np.array([i]) for i in range(1000)])
pi = tf.Variable(2.)

model = tf.keras.Sequential(
        [
         tf.keras.layers.InputLayer(input_shape=(1,)),
         tf.keras.layers.Dense(units = 1)
                ])


@tf.function
def compute_loss(model, x,pi):
    y_pred = model(x)
    return tf.keras.losses.MSE(x,y_pred) + pi**2

@tf.function
def compute_apply_gradients(model, x, optimizer):
  with tf.GradientTape() as tape:
    loss = compute_loss(model, x,pi)
  var_list = model.trainable_variables+[pi]
  gradients = tape.gradient(loss, var_list)
  optimizer.apply_gradients(zip(gradients, var_list))

#this doesnt work
for train_x in dataset_train:
    compute_apply_gradients(model, train_x, optimizer)


#######################################################################################
#instead of the above, this works and it computes 3 gradients, as expected ( 2 model parameters, 1 tf.variable).
with tf.GradientTape() as tape:
    loss = compute_loss(model, dataset_train,pi)
    gradients = tape.gradient(loss, model.trainable_variables+[pi])

все нормально при работе

        gradients = tape.gradient(loss, model.trainable_variables+[pi])

, который производит 3 градиента, но работает

for train_x in dataset_train:
    compute_apply_gradients(model, train_x, optimizer)

выдает следующую ошибку:

 raise e.ag_error_metadata.to_exception(e)

ValueError: in converted code:

    <ipython-input-65-750d9a39e3bf>:4 compute_apply_gradients  *
        loss = compute_loss(model, x,pi)
    /home/dan/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py:568 __call__
        result = self._call(*args, **kwds)
    <ipython-input-31-a9c6e2484e2a>:24 compute_loss  *
        y_pred = model(x)
    /home/dan/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py:737 __call__
        self.name)
    /home/dan/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/input_spec.py:193 assert_input_compatibility
        str(x.shape.as_list()))

    ValueError: Input 0 of layer sequential_4 is incompatible with the layer: : expected min_ndim=2, found ndim=1. Full shape received: [1]

и работает

  optimizer.apply_gradients(zip(gradients, var_list))

выдает следующую ошибку:

optimizer.apply_gradients(zip(gradients, var_list))

TypeError: apply_gradients() missing 1 required positional argument: 'grads_and_vars

Как это исправить? Спасибо!