python tensorflow dictionary eager-execution

Updating dictionary with tensorflow 2.1.0

I am facing a problem and I cannot seem to find the solution anywhere else, so I decided to post my question here (I have basic knowledge of tensorflow but quite new):

I wrote a simple code in python to illustrate what I want to do.

    import tensorflow as tf

    def generated_dict():
       graph = {'input': tf.Variable(2)}
       graph['layer_1'] = tf.square(graph['input'])
       graph['layer_2'] = tf.add(graph['input'], graph['layer_1'])
       return graph

    graph = generated_dict()
    print("boo = " + str(graph['layer_2']))

    graph['input'].assign(tf.constant(3))
    print("far = " + str(graph['layer_2']))

On this sample code, I would like tensorflow to update the whole dictionary when I assign a new input value by doing graph['input'].assign(tf.constant(3)) . Basically, right now I obtain

boo = tf.Tensor(6, shape=(), dtype=int32) # 2²+2
far = tf.Tensor(6, shape=(), dtype=int32) # 2²+2

which is normal because of eager execution of my code. However I would like the dictionary to update its values with my new input and to get :

boo = tf.Tensor(6, shape=(), dtype=int32) #2²+2
far = tf.Tensor(12, shape=(), dtype=int32) #3²+3

I have the feeling I should be using tf.function() but I am not sure how I should proceed with it. I tried graph = tf.function(generated_graph)() but I did not help. Any help will be greatly appreciated.

Solution

First of all, when trying to adapt TF1 code into TF2, I suggest to read the guide : Migrate your TensorFlow 1 code to TensorFlow 2.

TF2 changed the base design of tensorflow from running ops in a graph to eager execution. It means that writing code in TF2 is quite close to writing code in normal python : all the abstraction, graph creation, etc, is done under the hood.

Your complicated design does not need to exists in TF2, just write a simple python function. You can even use normal operators instead of tensorflow functions. There will be converted into tensorflow operators. Optionally, you can use the tf.function decorator for performances.

@tf.function
def my_graph(x):
     return x**2 + x

Now, if you want to feed data in that function, you just need to call it with a value :

>>> my_graph(tf.constant(3))
<tf.Tensor: shape=(), dtype=int32, numpy=12>
>>> my_graph(tf.constant(2))
<tf.Tensor: shape=(), dtype=int32, numpy=6>