AttributeError: module 'tensorflow_core.keras.backend' has no attribute 'track_variable'

[kechan]

I am using tf 2.0.0

Describe the current behavior
I have used a lambda layer in my model and declared 2 trainable vars inside the custom function. I saved and tried to load_model it back, and this seemed to be causing:

/tensorflow-2.0.0/python3.6/tensorflow_core/python/keras/layers/core.py in _variable_creator(self, next_creator, **kwargs)
    805     else:
    806       self._non_trainable_weights.append(var)
--> 807     K.track_variable(var)
    808     return var
    809 
AttributeError: module 'tensorflow_core.keras.backend' has no attribute 'track_variable'

Describe the expected behavior
No such error. Backend should have attribute 'track_variable' since it is code written in core.py. Although my sample code may still have something wrong with it, but code base should at least be self-consistent.

Code to reproduce the issue

def lambda_bug(x):
  ws = lambda: tf.ones(shape=(1, 1), dtype=tf.float32)

  w = tf.Variable(initial_value=ws, dtype=tf.float32, name='weights')
  b = tf.Variable(0., name='bias')

  return w*x + b

X = Input(shape=(1,))

out = Lambda(lambda_bug, name='lambda_bug')(X)
model = Model(inputs = X, outputs = out)
# model.summary() looks ok

model.save('buggy_model.h5')
model2 = load_model('buggy_model.h5')

[kechan]

Additional note: the model "works" in the sense if i provide an input:

x = np.ones((10,1), dtype='float32') yhat = model(x)

tf 2.0 will eagerly output yhat tensor without error. The problem seems to be with serializing a model that has trainable vars declared inside a lambda layer.

Beside this issue, if someone spotted any API or coding mistake, please let me know. I may get away with this 'track_variable' business. However, i do think this missing attribute has to be a bug, although I don't know what it is for.

[kechan]

I tested with save_weights(...) and load_weights(...) and I don't encounter this error. (so i am good as far as immediate work is concerned).