Updating dictionary with tensorflow 2.1.0

  dictionary, eager-execution, python, tensorflow

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']))

    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.

Source: Python Questions