Probably you could have some impact with:

  

      github.com/tensorflow/tensorflow
  

  
    
  
    
  

  
    
      Lifting variable on retrace 2
    

    
      tensorflow:master ← bhack:patch-18
    

    
      
        opened 
        
          
        
        May 19, 2021
      
      

      
        
          
          bhack
        
      

      
        
          +68
          -3
        
      
    
  


  
    Explore the effect on tests to fix: https://github.com/tensorflow/tensorflow/iss…ues/27120
  

  

  
    
    
  

  



  

      github.com/tensorflow/tensorflow
  

  
    
  
	  
  

  
    
      tf.function-decorated function tried to create variables on non-first call 3
    

    
      
        opened 
        
          
        
        Mar 26, 2019
      
      


      
        
          
          ericpts
        
      
    

    
        TF 2.0
        comp:ops
        type:bug
    
  


  
    **System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04…): Debian Testing
- TensorFlow installed from (source or binary): Binary
- TensorFlow version (use command below): 2.0.0.dev20190227
- Python version: 3.7


**Describe the current behavior**

A function which correctly works when in eager mode does not work anymore when annotated with `tf.function`.

In particular, it complains about `ValueError: tf.function-decorated function tried to create variables on non-first call.`, even though the function is always called with different parameters.

This is a continuation of https://github.com/tensorflow/tensorflow/issues/26812#issuecomment-475600836.

**Describe the expected behavior**

The `apply_gradients_once()` function should work even when annotated with `tf.function`.

**Code to reproduce the issue**

```python3
import tensorflow as tf
import numpy as np


fast_optimizer = tf.keras.optimizers.Adam(
        learning_rate=1e-3)

slow_optimizer = tf.keras.optimizers.Adam(
        learning_rate=1e-3 * 1e-9)


@tf.function
def apply_gradients_once(optimizer, grads, vars):
    grads = [grads]
    optimizer.apply_gradients(zip(grads, vars))


def apply_grads(use_fast, grads_per_model, vars):
    for i in range(2):
        if use_fast[i]:
            apply_gradients_once(fast_optimizer, grads_per_model[i], vars[i])
        else:
            apply_gradients_once(slow_optimizer, grads_per_model[i], vars[i])


def compute_loss(w, x, y):
    r = (w * x - y)**2
    r = tf.math.reduce_mean(r)
    return r

def compute_gradients(model):
    with tf.GradientTape() as tape:
        tape.watch(model)
        loss = compute_loss(model, x, y)
    grads = tape.gradient(loss, model)
    return grads


w = [
    tf.Variable(0.0),
    tf.Variable(1.0)]

x = np.array([1, 2, 3])
y = np.array([1, 2, 3])

vars = []
grads = []
for i in range(2):
    vars.append([w[i]])
    grads.append(compute_gradients(w[i]))

apply_grads([True, False], grads, vars)
```


**Other info / logs**

Error log:

```
Traceback (most recent call last):
  File "main.py", line 52, in <module>
    apply_grads([True, False], grads, vars)
  File "main.py", line 23, in apply_grads
    apply_gradients_once(slow_optimizer, grads_per_model[i], vars[i])
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 414, in __call__
    return self._stateless_fn(*args, **kwds)  # pylint: disable=not-callable
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 1254, in __call__
    graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 1577, in _maybe_define_function
    args, kwargs, override_flat_arg_shapes=relaxed_arg_shapes)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py", line 1479, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 685, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 317, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py", line 677, in wrapper
    ), args, kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/impl/api.py", line 392, in converted_call
    result = converted_f(*effective_args, **kwargs)
  File "/tmp/tmpr2ti5o1e.py", line 4, in tf__apply_gradients_once
    ag__.converted_call('apply_gradients', optimizer, ag__.ConversionOptions(recursive=True, verbose=0, strip_decorators=(tf.function, defun, ag__.convert, ag__.do_not_convert, ag__.converted_call), force_conversion=False, optional_features=(), internal_convert_user_code=True), (zip(grads, vars),), {})
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/impl/api.py", line 267, in converted_call
    return _call_unconverted(f, args, kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/autograph/impl/api.py", line 188, in _call_unconverted
    return f(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py", line 399, in apply_gradients
    self._create_hypers()
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py", line 558, in _create_hypers
    aggregation=tf_variables.VariableAggregation.ONLY_FIRST_REPLICA)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py", line 727, in add_weight
    aggregation=aggregation)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/training/tracking/base.py", line 622, in _add_variable_with_custom_getter
    **kwargs_for_getter)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer_utils.py", line 152, in make_variable
    aggregation=aggregation)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/variables.py", line 212, in __call__
    return cls._variable_v1_call(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/variables.py", line 175, in _variable_v1_call
    aggregation=aggregation)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/variables.py", line 58, in getter
    return captured_getter(captured_previous, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py", line 375, in invalid_creator_scope
    "tf.function-decorated function tried to create "
ValueError: tf.function-decorated function tried to create variables on non-first call.
```
  

  

  
    
    
  

  


But I’ve not tested your specific case with this new flag.