With TF 2.0 preview, "old-style" variational layers error with "Inputs to eager execution function cannot be Keras symbolic tensors"

[skeydan]

Hi,

I have unit tests for Convolution1DReparameterization , Convolution1DReparameterization etc. that basically look like this

x = tf.ones(shape = [150,1])
y = tf.ones(shape = [150])

model = tf.keras.Sequential(
  [tfpl.DenseReparameterization(
      units = 512,
      activation = "relu"),
   tfpl.DenseReparameterization(
      units = 1)])

model.compile(optimizer = 'adam', loss = "mse")
model.fit(x, y, steps_per_epoch = 1)

These run fine with TF 1, but with TF 2 preview I get

_SymbolicException: Inputs to eager execution function cannot be Keras symbolic tensors, but found [<tf.Tensor 'sequential/dense_reparameterization/divergence_kernel:0' shape=() dtype=float32>, <tf.Tensor 'sequential/dense_reparameterization_1/divergence_kernel:0' shape=() dtype=float32>]

Would you happen to know what is the cause here?
Are these "old-style" variational layers supposed to run with TF 2?

Many thanks in advance!!

[jburnim]

These layers should be working in TF2.

It looks like there is an incompatibility with these layers and a recent change in Keras. We will investigate.

If you add experimental_run_tf_function=False as an argument to model.compile, do the errors go away?

[skeydan]

Thanks, and

If you add experimental_run_tf_function=False as an argument to model.compile, do the errors go away?

yes they do!

[vgoklani]

any reason why we have to use "experimental_run_tf_function=False". I need to use this too.

[skeydan]

In the TF 2.0-rc0 release notes I saw this

Switched Keras fit/evaluate/predict execution to use only a single unified path by default unless eager execution has been explicitly disabled, regardless of input type. This unified path places an eager-friendly training step inside of a tf.function. With this 1. All input types are converted to Dataset. 2. The path assumes there is always a distribution strategy. when distribution strategy is not specified the path uses a no-op distribution strategy. 3. The training step is wrapped in tf.function unless run_eagerly=True is set in compile. The single path execution code does not yet support all use cases. We fallback to the existing v1 execution paths if your model contains the following: 1. sample_weight_mode in compile 2. weighted_metrics in compile 3. v1 optimizer 4. target tensors in compile. If you are experiencing any issues because of this change, please inform us (file an issue) about your use case and you can unblock yourself by setting experimental_run_tf_function=False in compile meanwhile. We have seen couple of use cases where the model usage pattern is not as expected and would not work with this change. 1. output tensors of one layer is used in the constructor of another. 2. symbolic tensors outside the scope of the model are used in custom loss functions. The flag can be disabled for these cases and ideally the usage pattern will need to be fixed.

So guessing (haven't checked) it might be "output tensors of one layer is used in the constructor of another."

However @jburnim as far as I can see these layers don't really work with TF 2, do they (assuming eager mode)?
I tried Keras-style fit which seemed to run fine syntactically, but didn't accumulate the losses (and furthermore, doesn't allow for scaling the KL), as well as GradientTape which gave an error (for all this see mails to the list: https://groups.google.com/a/tensorflow.org/forum/#!topic/tfprobability/axQpJqs0oHc ).

Is there a way to use these layers in TF 2, and if so, could you please indicate it? Otherwise, are there plans to create V2 versions of these layers?

Thanks!

[tomtang110]

experimental_run_tf_function=False

but If I don't use model.compile, how do I do?
`class PolicyEstimator():

def __init__(self, lr=1e-2, name='policy_estimator'):
    self.state = tf.keras.Input(shape=(400,), name='state', dtype=tf.float32)
    self.mu = layers.Dense(1, activation=None,kernel_initializer='random_normal')(self.state)
    self.mu = tf.squeeze(self.mu)
    self.sigma = layers.Dense(1, activation=None,kernel_initializer='random_normal')(self.state)
    self.sigma = tf.squeeze(self.sigma)
    self.sigma = tf.keras.activations.softplus(self.sigma) + 1e-5
    self.normal_dist = tfp.distributions.Normal(self.mu, self.sigma)
    self.action = self.normal_dist._sample_n(1)
    self.action = tf.clip_by_value(self.action, env.action_space.low[0], env.action_space.high[0])
    self.model = tf.keras.Model(inputs=self.state, outputs=self.action)
    self.optimizer = tf.keras.optimizers.Adam(lr)

def predict(self, state):
    return self.model(state)

def update(self, state, target, action):
    with tf.GradientTape() as tape:
        predictions = self.model(state)[0]
        loss = -self.normal_dist.log_prob(predictions) # this would produce a non-eager tensor
        loss -= 1e-1 * self.normal_dist.entropy()
        print('policy: ', loss)
        grads = tape.gradient(loss, self.model.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables))
    return loss`

As you can be seen, the loss in update function is produced by self.normal_dist.log_prob, which would produce the error-"Inputs to eager execution function cannot be Keras symbolic tensors,but found [<tf.Tensor 'StridedSliceGrad:0' shape=(1,) dtype=float32>]".
I did not use compile function so how can I solve the problem?

[Crispy13]

These layers should be working in TF2.

It looks like there is an incompatibility with these layers and a recent change in Keras. We will investigate.

If you add experimental_run_tf_function=False as an argument to model.compile, do the errors go away?

I added it as an argument to model.compile, but didn't work.

[felixvelariusbos]

Just an added update here for version: This issue is still in TensorFlow 2.1, but the provided workaround (experimental_run_tf_function=False) still works

[William-YanHua]

These layers should be working in TF2.

It looks like there is an incompatibility with these layers and a recent change in Keras. We will investigate.

If you add experimental_run_tf_function=False as an argument to model.compile, do the errors go away?

I add this argument to 'model.compile', it can be compiled successfully, but when it execution with ModelCheckpoint and EarlyStopping, it will cause the same problem while saving the checkpoint.