Bug using optimizer with tf.function

[ShHsLin]

I encounter a problem regarding using optimizer with tf.function. It seems that the optimizer would not work under tf.function decoration. I attach a simple example showing the behaviour.

If one remove the @tf.function in front of one_step, the code works. Also one use the normal optimizer, i.e. tf.optimizers, the code would works.

I wonder if there is anything underlying make the autograph tracing broken.

import numpy as np
import tensorflow as tf
import QGOpt as qgo

m = qgo.manifolds.StiefelManifold()
opt = qgo.optimizers.RAdam(m, 0.5, ams=True) # <- with this and the code does not work
# opt = tf.optimizers.Adam(0.5) # <- code works fine

uc = m.random((3, 3), dtype=tf.complex128)
u = tf.Variable(qgo.manifolds.complex_to_real(uc))

@tf.function
def loss_simple(uc):
    return tf.math.real(tf.reduce_sum(uc))

@tf.function   # <- removing this and the code works fine
def one_step(u, opt):
    with tf.GradientTape() as tape:
        uc = tf.complex(u[...,0], u[...,1])
        loss = loss_simple(uc)

    grads = tape.gradient(loss, u)
    opt.apply_gradients(zip([grads], [u]))
    return loss

print(u)
print(one_step(u, opt))
print(u)

would show

    TypeError: Can not convert a NoneType into a Tensor or Operation.

[LuchnikovI]

Thanks for your finding, I will get back to this problem soon.

[LuchnikovI]

The problem appears when one calls method _distributed_apply, which is not modified in QGOpt. There is a short guide on how to implement subclasses of the class Optimizer here https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/optimizer.py, and it is not clear for now what to do with _distributed_apply.