Tensorflow UnimplementedError

Я реализовал в Tensorflow следующий пример:

import tensorflow as tf
import numpy as np


def loss_function(values, a, b):
    N = values.shape[0]
    i = tf.constant(0)
    values_array = tf.get_variable(
        "values", values.shape, initializer=tf.constant_initializer(values))
    result = tf.constant(0, dtype=tf.float32)

    def body1(i):

        op2 = tf.assign(values_array[i, 0],
                        a + b)  # Here is where it should be updated. The value being assigned is actually calculated from variable a and b.

        with tf.control_dependencies([op2]):
            return tf.identity(i + 1)

    def condition1(i): return tf.less(i, N)
    i = tf.while_loop(condition1, body1, [i])

    with tf.control_dependencies([i]):
        op1 = tf.assign(values_array[0, 0],
                        999.0)  # Here is where it should be updated

        # The final cost is calculated based on the entire values_array
        with tf.control_dependencies([op1]):
            result = result + tf.reduce_mean(values_array)
            return tf.identity(result)


# The parameters we want to calculate in the end
a = tf.Variable(tf.random_uniform([1], 0, 700), name='a')
b = tf.Variable(tf.random_uniform([1], -700, 700), name='b')

values = np.ones([2, 4], dtype=np.float32)

# cost function
cost_function = loss_function(values, a, b)

# training algorithm
optimizer = tf.train.MomentumOptimizer(
    0.1, momentum=0.9).minimize(cost_function)

# initializing the variables
init = tf.global_variables_initializer()

# starting the session session
sess = tf.Session()
sess.run(init)

training_cost = sess.run([cost_function])
_ = sess.run([optimizer])

print tf.get_collection(
    tf.GraphKeys.GLOBAL_VARIABLES, scope = "values")[0].eval(session=sess)

В целом, я хочу, чтобы функция стоимости вычисляла временный 2D-массив на основе входного numpy 2D-массива и параметров a и b. Затем окончательная стоимость рассчитывается на основе временного 2D-массива. Однако это вызывает исключение Exception has occurred: UnimplementedError UnimplementedError().

Любая помощь?

Спасибо!