Loops, Conditionals & AutoGraph: Writing Graph-Friendly TensorFlow Code

Table of Contents AutoGraph transformations Conditionals Loops Limitations Executing Python side effects All outputs of a tf.function must be return values Recursive tf.fuctions are not supported. Known issues Depending on Python global and free variables Depending on Python objects Creating tf.Variables AutoGraph transformations Conditionals Loops Conditionals Loops Conditionals Loops Limitations Executing Python side effects All outputs of a tf.function must be return values Recursive tf.fuctions are not supported. Executing Python side effects All outputs of a tf.function must be return values Recursive tf.fuctions are not supported. Executing Python side effects All outputs of a tf.function must be return values tf.function Recursive tf.fuctions are not supported. Known issues Depending on Python global and free variables Depending on Python objects Creating tf.Variables Depending on Python global and free variables Depending on Python objects Creating tf.Variables Depending on Python global and free variables Depending on Python objects Creating tf.Variables AutoGraph transformations AutoGraph is a library that is on by default in tf.function, and transforms a subset of Python eager code into graph-compatible TensorFlow ops. This includes control flow like if, for, while. tf.function if for while TensorFlow ops like tf.cond and tf.while_loop continue to work, but control flow is often easier to write and understand when written in Python. tf.cond tf.while_loop # A simple loop @tf.function def f(x): while tf.reduce_sum(x) > 1: tf.print(x) x = tf.tanh(x) return x f(tf.random.uniform([5])) # A simple loop @tf.function def f(x): while tf.reduce_sum(x) > 1: tf.print(x) x = tf.tanh(x) return x f(tf.random.uniform([5])) [0.722626925 0.640327692 0.725044 0.904435039 0.868018746] [0.61853379 0.565122604 0.620023966 0.718450606 0.700366139] [0.550106347 0.511768281 0.551144719 0.615948677 0.604600191] [0.500599921 0.471321791 0.501377642 0.548301 0.540314913] [0.462588847 0.439266682 0.463199914 0.499245733 0.493226349] [0.432191819 0.413036436 0.432688653 0.461523771 0.456773371] [0.407151431 0.391047835 0.407565802 0.431325316 0.427450746] [0.386051297 0.372263193 0.386403859 0.406428277 0.403188676] [0.367951065 0.355969697 0.368255854 0.38543576 0.382673979] [0.352198243 0.341659099 0.352465183 0.367418766 0.365027398] [0.338323593 0.328957736 0.33856 0.351731867 0.349634588] [0.325979948 0.317583948 0.326191217 0.337910533 0.336051434] [0.314903945 0.307320684 0.315094262 0.325610697 0.323947728] [0.304891765 0.297997624 0.30506441 0.314571291 0.313072115] [0.295782804 0.289479077 0.29594034 0.304590017 0.303229302] [0.287448555 0.281655282 0.287593067 0.295507431 0.294265062] [0.279784769 0.274436355 0.279917955 0.287195921 0.286055595] [0.272705853 0.267748028 0.272829145 0.279551893 0.278500348] [0.266140789 0.261528105 0.266255379 0.272490293 0.271516532] [0.26003018 0.255724251 0.260137022 0.265940517 0.265035421] [0.254323781 0.250291914 0.254423678 0.259843439 0.258999288] [0.248978764 0.245193034 0.249072418 0.25414905 0.253359258] [0.243958414 0.240394741 0.244046524 0.248814836 0.248073786] [0.239231125 0.235868543 0.239314198 0.243804231 0.24310714] [0.234769359 0.231589615 0.234847859 0.239085764 0.238428399] [0.230549142 0.227536201 0.230623439 0.234632015 0.234010741] [0.226549357 0.223689109 0.22661984 0.23041907 0.229830697] [0.222751439 0.220031396 0.222818434 0.226425976 0.225867674] [0.21913895 0.216548 0.219202697 0.222634196 0.222103462] [0.215697214 0.213225439 0.215757981 0.219027311 0.218521982] [0.212413162 0.210051686 0.212471202 0.215590775 0.215108871] [0.209275112 0.207015961 0.209330618 0.212311521 0.211851314] [0.206272557 0.204108506 0.206325665 0.209177911 0.20873782] [0.203395993 0.201320544 0.203446865 0.206179485 0.20575805] [0.200636819 0.198644072 0.200685605 0.203306749 0.202902704] [0.722626925 0.640327692 0.725044 0.904435039 0.868018746] [0.61853379 0.565122604 0.620023966 0.718450606 0.700366139] [0.550106347 0.511768281 0.551144719 0.615948677 0.604600191] [0.500599921 0.471321791 0.501377642 0.548301 0.540314913] [0.462588847 0.439266682 0.463199914 0.499245733 0.493226349] [0.432191819 0.413036436 0.432688653 0.461523771 0.456773371] [0.407151431 0.391047835 0.407565802 0.431325316 0.427450746] [0.386051297 0.372263193 0.386403859 0.406428277 0.403188676] [0.367951065 0.355969697 0.368255854 0.38543576 0.382673979] [0.352198243 0.341659099 0.352465183 0.367418766 0.365027398] [0.338323593 0.328957736 0.33856 0.351731867 0.349634588] [0.325979948 0.317583948 0.326191217 0.337910533 0.336051434] [0.314903945 0.307320684 0.315094262 0.325610697 0.323947728] [0.304891765 0.297997624 0.30506441 0.314571291 0.313072115] [0.295782804 0.289479077 0.29594034 0.304590017 0.303229302] [0.287448555 0.281655282 0.287593067 0.295507431 0.294265062] [0.279784769 0.274436355 0.279917955 0.287195921 0.286055595] [0.272705853 0.267748028 0.272829145 0.279551893 0.278500348] [0.266140789 0.261528105 0.266255379 0.272490293 0.271516532] [0.26003018 0.255724251 0.260137022 0.265940517 0.265035421] [0.254323781 0.250291914 0.254423678 0.259843439 0.258999288] [0.248978764 0.245193034 0.249072418 0.25414905 0.253359258] [0.243958414 0.240394741 0.244046524 0.248814836 0.248073786] [0.239231125 0.235868543 0.239314198 0.243804231 0.24310714] [0.234769359 0.231589615 0.234847859 0.239085764 0.238428399] [0.230549142 0.227536201 0.230623439 0.234632015 0.234010741] [0.226549357 0.223689109 0.22661984 0.23041907 0.229830697] [0.222751439 0.220031396 0.222818434 0.226425976 0.225867674] [0.21913895 0.216548 0.219202697 0.222634196 0.222103462] [0.215697214 0.213225439 0.215757981 0.219027311 0.218521982] [0.212413162 0.210051686 0.212471202 0.215590775 0.215108871] [0.209275112 0.207015961 0.209330618 0.212311521 0.211851314] [0.206272557 0.204108506 0.206325665 0.209177911 0.20873782] [0.203395993 0.201320544 0.203446865 0.206179485 0.20575805] [0.200636819 0.198644072 0.200685605 0.203306749 0.202902704] If you're curious you can inspect the code AutoGraph generates. print(tf.autograph.to_code(f.python_function)) print(tf.autograph.to_code(f.python_function)) def tf__f(x): with ag__.FunctionScope('f', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope: do_return = False retval_ = ag__.UndefinedReturnValue() def get_state(): return (x,) def set_state(vars_): nonlocal x (x,) = vars_ def loop_body(): nonlocal x ag__.converted_call(ag__.ld(tf).print, (ag__.ld(x),), None, fscope) x = ag__.converted_call(ag__.ld(tf).tanh, (ag__.ld(x),), None, fscope) def loop_test(): return ag__.converted_call(ag__.ld(tf).reduce_sum, (ag__.ld(x),), None, fscope) > 1 ag__.while_stmt(loop_test, loop_body, get_state, set_state, ('x',), {}) try: do_return = True retval_ = ag__.ld(x) except: do_return = False raise return fscope.ret(retval_, do_return) def tf__f(x): with ag__.FunctionScope('f', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope: do_return = False retval_ = ag__.UndefinedReturnValue() def get_state(): return (x,) def set_state(vars_): nonlocal x (x,) = vars_ def loop_body(): nonlocal x ag__.converted_call(ag__.ld(tf).print, (ag__.ld(x),), None, fscope) x = ag__.converted_call(ag__.ld(tf).tanh, (ag__.ld(x),), None, fscope) def loop_test(): return ag__.converted_call(ag__.ld(tf).reduce_sum, (ag__.ld(x),), None, fscope) > 1 ag__.while_stmt(loop_test, loop_body, get_state, set_state, ('x',), {}) try: do_return = True retval_ = ag__.ld(x) except: do_return = False raise return fscope.ret(retval_, do_return) Conditionals AutoGraph will convert some if statements into the equivalent tf.cond calls. This substitution is made if is a Tensor. Otherwise, the if statement is executed as a Python conditional. if tf.cond if A Python conditional executes during tracing, so exactly one branch of the conditional will be added to the graph. Without AutoGraph, this traced graph would be unable to take the alternate branch if there is data-dependent control flow. tf.cond traces and adds both branches of the conditional to the graph, dynamically selecting a branch at execution time. Tracing can have unintended side effects; check out AutoGraph tracing effects for more information. tf.cond AutoGraph tracing effects tf.function def fizzbuzz(n): for i in tf.range(1, n + 1): print('Tracing for loop') if i % 15 == 0: print('Tracing fizzbuzz branch') tf.print('fizzbuzz') elif i % 3 == 0: print('Tracing fizz branch') tf.print('fizz') elif i % 5 == 0: print('Tracing buzz branch') tf.print('buzz') else: print('Tracing default branch') tf.print(i) fizzbuzz(tf.constant(5)) fizzbuzz(tf.constant(20)) tf.function def fizzbuzz(n): for i in tf.range(1, n + 1): print('Tracing for loop') if i % 15 == 0: print('Tracing fizzbuzz branch') tf.print('fizzbuzz') elif i % 3 == 0: print('Tracing fizz branch') tf.print('fizz') elif i % 5 == 0: print('Tracing buzz branch') tf.print('buzz') else: print('Tracing default branch') tf.print(i) fizzbuzz(tf.constant(5)) fizzbuzz(tf.constant(20)) Tracing for loop Tracing fizzbuzz branch Tracing fizz branch Tracing buzz branch Tracing default branch 1 2 fizz 4 buzz 1 2 fizz 4 buzz fizz 7 8 fizz buzz 11 fizz 13 14 fizzbuzz 16 17 fizz 19 buzz Tracing for loop Tracing fizzbuzz branch Tracing fizz branch Tracing buzz branch Tracing default branch 1 2 fizz 4 buzz 1 2 fizz 4 buzz fizz 7 8 fizz buzz 11 fizz 13 14 fizzbuzz 16 17 fizz 19 buzz See the reference documentation for additional restrictions on AutoGraph-converted if statements. reference documentation Loops AutoGraph will convert some for and while statements into the equivalent TensorFlow looping ops, like tf.while_loop. If not converted, the for or while loop is executed as a Python loop. for while tf.while_loop for while This substitution is made in the following situations: for x in y: if y is a Tensor, convert to tf.while_loop. In the special case where y is a tf.data.Dataset, a combination of tf.data.Dataset ops are generated. while : if is a Tensor, convert to tf.while_loop. for x in y: if y is a Tensor, convert to tf.while_loop. In the special case where y is a tf.data.Dataset, a combination of tf.data.Dataset ops are generated. for x in y y tf.while_loop y tf.data.Dataset tf.data.Dataset while : if is a Tensor, convert to tf.while_loop. while tf.while_loop A Python loop executes during tracing, adding additional ops to the tf.Graph for every iteration of the loop. tf.Graph A TensorFlow loop traces the body of the loop, and dynamically selects how many iterations to run at execution time. The loop body only appears once in the generated tf.Graph. tf.Graph See the reference documentation for additional restrictions on AutoGraph-converted for and while statements. reference documentation for while Looping over Python data A common pitfall is to loop over Python/NumPy data within a tf.function. This loop will execute during the tracing process, adding a copy of your model to the tf.Graph for each iteration of the loop. tf.function tf.Graph If you want to wrap the entire training loop in tf.function, the safest way to do this is to wrap your data as a tf.data.Dataset so that AutoGraph will dynamically unroll the training loop. tf.function tf.data.Dataset def measure_graph_size(f, *args): g = f.get_concrete_function(*args).graph print("{}({}) contains {} nodes in its graph".format( f.__name__, ', '.join(map(str, args)), len(g.as_graph_def().node))) @tf.function def train(dataset): loss = tf.constant(0) for x, y in dataset: loss += tf.abs(y - x) # Some dummy computation. return loss small_data = [(1, 1)] * 3 big_data = [(1, 1)] * 10 measure_graph_size(train, small_data) measure_graph_size(train, big_data) measure_graph_size(train, tf.data.Dataset.from_generator( lambda: small_data, (tf.int32, tf.int32))) measure_graph_size(train, tf.data.Dataset.from_generator( lambda: big_data, (tf.int32, tf.int32))) def measure_graph_size(f, *args): g = f.get_concrete_function(*args).graph print("{}({}) contains {} nodes in its graph".format( f.__name__, ', '.join(map(str, args)), len(g.as_graph_def().node))) @tf.function def train(dataset): loss = tf.constant(0) for x, y in dataset: loss += tf.abs(y - x) # Some dummy computation. return loss small_data = [(1, 1)] * 3 big_data = [(1, 1)] * 10 measure_graph_size(train, small_data) measure_graph_size(train, big_data) measure_graph_size(train, tf.data.Dataset.from_generator( lambda: small_data, (tf.int32, tf.int32))) measure_graph_size(train, tf.data.Dataset.from_generator( lambda: big_data, (tf.int32, tf.int32))) train([(1, 1), (1, 1), (1, 1)]) contains 11 nodes in its graph train([(1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1)]) contains 32 nodes in its graph train( , dtype=tf.int32, name=None), TensorSpec(shape= , dtype=tf.int32, name=None))>) contains 6 nodes in its graph train( , dtype=tf.int32, name=None), TensorSpec(shape= , dtype=tf.int32, name=None))>) contains 6 nodes in its graph train([(1, 1), (1, 1), (1, 1)]) contains 11 nodes in its graph train([(1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1)]) contains 32 nodes in its graph train( , dtype=tf.int32, name=None), TensorSpec(shape= , dtype=tf.int32, name=None))>) contains 6 nodes in its graph train( , dtype=tf.int32, name=None), TensorSpec(shape= , dtype=tf.int32, name=None))>) contains 6 nodes in its graph When wrapping Python/NumPy data in a Dataset, be mindful of tf.data.Dataset.from_generator versus tf.data.Dataset.from_tensor_slices. The former will keep the data in Python and fetch it via tf.py_function which can have performance implications, whereas the latter will bundle a copy of the data as one large tf.constant() node in the graph, which can have memory implications. tf.data.Dataset.from_generator tf.data.Dataset.from_tensor_slices tf.py_function tf.constant() Reading data from files via TFRecordDataset, CsvDataset, etc. is the most effective way to consume data, as then TensorFlow itself can manage the asynchronous loading and prefetching of data, without having to involve Python. To learn more, see the tf.data: Build TensorFlow input pipelines guide. TFRecordDataset CsvDataset tf.data Accumulating values in a loop A common pattern is to accumulate intermediate values from a loop. Normally, this is accomplished by appending to a Python list or adding entries to a Python dictionary. However, as these are Python side effects, they will not work as expected in a dynamically unrolled loop. Use tf.TensorArray to accumulate results from a dynamically unrolled loop. tf.TensorArray batch_size = 2 seq_len = 3 feature_size = 4 def rnn_step(inp, state): return inp + state @tf.function def dynamic_rnn(rnn_step, input_data, initial_state): # [batch, time, features] -> [time, batch, features] input_data = tf.transpose(input_data, [1, 0, 2]) max_seq_len = input_data.shape[0] states = tf.TensorArray(tf.float32, size=max_seq_len) state = initial_state for i in tf.range(max_seq_len): state = rnn_step(input_data[i], state) states = states.write(i, state) return tf.transpose(states.stack(), [1, 0, 2]) dynamic_rnn(rnn_step, tf.random.uniform([batch_size, seq_len, feature_size]), tf.zeros([batch_size, feature_size])) batch_size = 2 seq_len = 3 feature_size = 4 def rnn_step(inp, state): return inp + state @tf.function def dynamic_rnn(rnn_step, input_data, initial_state): # [batch, time, features] -> [time, batch, features] input_data = tf.transpose(input_data, [1, 0, 2]) max_seq_len = input_data.shape[0] states = tf.TensorArray(tf.float32, size=max_seq_len) state = initial_state for i in tf.range(max_seq_len): state = rnn_step(input_data[i], state) states = states.write(i, state) return tf.transpose(states.stack(), [1, 0, 2]) dynamic_rnn(rnn_step, tf.random.uniform([batch_size, seq_len, feature_size]), tf.zeros([batch_size, feature_size])) <tf.Tensor: shape=(2, 3, 4), dtype=float32, numpy= array([[[0.98815036, 0.8358947 , 0.15233278, 0.58257985], [1.7802314 , 1.215749 , 0.6186948 , 0.9416343 ], [2.1005788 , 1.2919371 , 1.1675987 , 1.4443643 ]], [[0.751495 , 0.8949536 , 0.16761959, 0.45424747], [0.9617816 , 1.7412133 , 0.37147725, 0.7925167 ], [1.655664 , 1.9362986 , 1.1732976 , 1.12577 ]]], dtype=float32)> <tf.Tensor: shape=(2, 3, 4), dtype=float32, numpy= array([[[0.98815036, 0.8358947 , 0.15233278, 0.58257985], [1.7802314 , 1.215749 , 0.6186948 , 0.9416343 ], [2.1005788 , 1.2919371 , 1.1675987 , 1.4443643 ]], [[0.751495 , 0.8949536 , 0.16761959, 0.45424747], [0.9617816 , 1.7412133 , 0.37147725, 0.7925167 ], [1.655664 , 1.9362986 , 1.1732976 , 1.12577 ]]], dtype=float32)> Limitations tf.function has a few limitations by design that you should be aware of when converting a Python function to a tf.function. tf.function tf.function Executing Python side effects Side effects, like printing, appending to lists, and mutating globals, can behave unexpectedly inside a tf.function, sometimes executing twice or not all. They only happen the first time you call a tf.function with a set of inputs. Afterwards, the traced tf.Graph is reexecuted, without executing the Python code. tf.function tf.function tf.Graph The general rule of thumb is to avoid relying on Python side effects in your logic and only use them to debug your traces. Otherwise, TensorFlow APIs like tf.data, tf.print, tf.summary, tf.Variable.assign, and tf.TensorArray are the best way to ensure your code will be executed by the TensorFlow runtime with each call. tf.data tf.print tf.summary tf.Variable.assign tf.TensorArray @tf.function def f(x): print("Traced with", x) tf.print("Executed with", x) f(1) f(1) f(2) @tf.function def f(x): print("Traced with", x) tf.print("Executed with", x) f(1) f(1) f(2) Traced with 1 Executed with 1 Executed with 1 Traced with 2 Executed with 2 Traced with 1 Executed with 1 Executed with 1 Traced with 2 Executed with 2 If you would like to execute Python code during each invocation of a tf.function, tf. py_function is an exit hatch. The drawbacks of tf.py_function are that it's not portable or particularly performant, cannot be saved with SavedModel, and does not work well in distributed (multi-GPU, TPU) setups. Also, since tf.py_function has to be wired into the graph, it casts all inputs/outputs to tensors. tf.function tf. py_function tf.py_function SavedModel tf.py_function @tf.py_function(Tout=tf.float32) def py_plus(x, y): print('Executing eagerly.') return x + y @tf.function def tf_wrapper(x, y): print('Tracing.') return py_plus(x, y) @tf.py_function(Tout=tf.float32) def py_plus(x, y): print('Executing eagerly.') return x + y @tf.function def tf_wrapper(x, y): print('Tracing.') return py_plus(x, y) The tf.function will trace the first time: tf.function tf_wrapper(tf.constant(1.0), tf.constant(2.0)).numpy() tf_wrapper(tf.constant(1.0), tf.constant(2.0)).numpy() Tracing. Executing eagerly. 3.0 Tracing. Executing eagerly. 3.0 But the tf.py_function inside executes eagerly every time: tf.py_function tf_wrapper(tf.constant(1.0), tf.constant(2.0)).numpy() tf_wrapper(tf.constant(1.0), tf.constant(2.0)).numpy() Executing eagerly. 3.0 Executing eagerly. 3.0 Changing Python global and free variables Changing Python global and free variables counts as a Python side effect, so it only happens during tracing. free variables external_list = [] @tf.function def side_effect(x): print('Python side effect') external_list.append(x) side_effect(1) side_effect(1) side_effect(1) # The list append only happened once! assert len(external_list) == 1 external_list = [] @tf.function def side_effect(x): print('Python side effect') external_list.append(x) side_effect(1) side_effect(1) side_effect(1) # The list append only happened once! assert len(external_list) == 1 Python side effect Python side effect Sometimes unexpected behaviors are very hard to notice. In the example below, the counter is intended to safeguard the increment of a variable. However because it is a python integer and not a TensorFlow object, it's value is captured during the first trace. When the tf.function is used, the assign_add will be recorded unconditionally in the underlying graph. Therefore v will increase by 1, every time the tf.function is called. This issue is common among users that try to migrate their Graph-mode Tensorflow code to Tensorflow 2 using tf.function decorators, when python side-effects (the counter in the example) are used to determine what ops to run (assign_add in the example). Usually, users realize this only after seeing suspicious numerical results, or significantly lower performance than expected (e.g. if the guarded operation is very costly). counter tf.function assign_add v tf.function tf.function counter assign_add class Model(tf.Module): def __init__(self): self.v = tf.Variable(0) self.counter = 0 @tf.function def __call__(self): if self.counter == 0: # A python side-effect self.counter += 1 self.v.assign_add(1) return self.v m = Model() for n in range(3): print(m().numpy()) # prints 1, 2, 3 class Model(tf.Module): def __init__(self): self.v = tf.Variable(0) self.counter = 0 @tf.function def __call__(self): if self.counter == 0: # A python side-effect self.counter += 1 self.v.assign_add(1) return self.v m = Model() for n in range(3): print(m().numpy()) # prints 1, 2, 3 1 2 3 1 2 3 A workaround to achieve the expected behavior is using tf.init_scope to lift the operations outside of the function graph. This ensures that the variable increment is only done once during tracing time. It should be noted init_scope has other side effects including cleared control flow and gradient tape. Sometimes the usage of init_scope can become too complex to manage realistically. tf.init_scope init_scope init_scope class Model(tf.Module): def __init__(self): self.v = tf.Variable(0) self.counter = 0 @tf.function def __call__(self): if self.counter == 0: # Lifts ops out of function-building graphs with tf.init_scope(): self.counter += 1 self.v.assign_add(1) return self.v m = Model() for n in range(3): print(m().numpy()) # prints 1, 1, 1 class Model(tf.Module): def __init__(self): self.v = tf.Variable(0) self.counter = 0 @tf.function def __call__(self): if self.counter == 0: # Lifts ops out of function-building graphs with tf.init_scope(): self.counter += 1 self.v.assign_add(1) return self.v m = Model() for n in range(3): print(m().numpy()) # prints 1, 1, 1 1 1 1 1 1 1 In summary, as a rule of thumb, you should avoid mutating python objects such as integers or containers like lists that live outside the tf.function. Instead, use arguments and TF objects. For example, the section "Accumulating values in a loop" has one example of how list-like operations can be implemented. tf.function "Accumulating values in a loop" You can, in some cases, capture and manipulate state if it is a tf.Variable. This is how the weights of Keras models are updated with repeated calls to the same ConcreteFunction. tf.Variable ConcreteFunction Using Python iterators and generators Many Python features, such as generators and iterators, rely on the Python runtime to keep track of state. In general, while these constructs work as expected in eager mode, they are examples of Python side effects and therefore only happen during tracing. @tf.function def buggy_consume_next(iterator): tf.print("Value:", next(iterator)) iterator = iter([1, 2, 3]) buggy_consume_next(iterator) # This reuses the first value from the iterator, rather than consuming the next value. buggy_consume_next(iterator) buggy_consume_next(iterator) @tf.function def buggy_consume_next(iterator): tf.print("Value:", next(iterator)) iterator = iter([1, 2, 3]) buggy_consume_next(iterator) # This reuses the first value from the iterator, rather than consuming the next value. buggy_consume_next(iterator) buggy_consume_next(iterator) Value: 1 Value: 1 Value: 1 Value: 1 Value: 1 Value: 1 Just like how TensorFlow has a specialized tf.TensorArray for list constructs, it has a specialized tf.data.Iterator for iteration constructs. See the section on AutoGraph transformations for an overview. Also, the tf.data API can help implement generator patterns: tf.TensorArray tf.data.Iterator AutoGraph transformations tf.data @tf.function def good_consume_next(iterator): # This is ok, iterator is a tf.data.Iterator tf.print("Value:", next(iterator)) ds = tf.data.Dataset.from_tensor_slices([1, 2, 3]) iterator = iter(ds) good_consume_next(iterator) good_consume_next(iterator) good_consume_next(iterator) @tf.function def good_consume_next(iterator): # This is ok, iterator is a tf.data.Iterator tf.print("Value:", next(iterator)) ds = tf.data.Dataset.from_tensor_slices([1, 2, 3]) iterator = iter(ds) good_consume_next(iterator) good_consume_next(iterator) good_consume_next(iterator) Value: 1 Value: 2 Value: 3 Value: 1 Value: 2 Value: 3 All outputs of a tf.function must be return values With the exception of tf.Variables, a tf.function must return all its outputs. Attempting to directly access any tensors from a function without going through return values causes "leaks". tf.Variable For example, the function below "leaks" the tensor a through the Python global x: a x x = None @tf.function def leaky_function(a): global x x = a + 1 # Bad - leaks local tensor return a + 2 correct_a = leaky_function(tf.constant(1)) print(correct_a.numpy()) # Good - value obtained from function's returns try: x.numpy() # Bad - tensor leaked from inside the function, cannot be used here except AttributeError as expected: print(expected) x = None @tf.function def leaky_function(a): global x x = a + 1 # Bad - leaks local tensor return a + 2 correct_a = leaky_function(tf.constant(1)) print(correct_a.numpy()) # Good - value obtained from function's returns try: x.numpy() # Bad - tensor leaked from inside the function, cannot be used here except AttributeError as expected: print(expected) 3 'SymbolicTensor' object has no attribute 'numpy' 3 'SymbolicTensor' object has no attribute 'numpy' This is true even if the leaked value is also returned: @tf.function def leaky_function(a): global x x = a + 1 # Bad - leaks local tensor return x # Good - uses local tensor correct_a = leaky_function(tf.constant(1)) print(correct_a.numpy()) # Good - value obtained from function's returns try: x.numpy() # Bad - tensor leaked from inside the function, cannot be used here except AttributeError as expected: print(expected) @tf.function def captures_leaked_tensor(b): b += x # Bad - `x` is leaked from `leaky_function` return b with assert_raises(TypeError): captures_leaked_tensor(tf.constant(2)) @tf.function def leaky_function(a): global x x = a + 1 # Bad - leaks local tensor return x # Good - uses local tensor correct_a = leaky_function(tf.constant(1)) print(correct_a.numpy()) # Good - value obtained from function's returns try: x.numpy() # Bad - tensor leaked from inside the function, cannot be used here except AttributeError as expected: print(expected) @tf.function def captures_leaked_tensor(b): b += x # Bad - `x` is leaked from `leaky_function` return b with assert_raises(TypeError): captures_leaked_tensor(tf.constant(2)) 2 'SymbolicTensor' object has no attribute 'numpy' Caught expected exception : Traceback (most recent call last): File "/tmpfs/tmp/ipykernel_167534/3551158538.py", line 8, in assert_raises yield File "/tmpfs/tmp/ipykernel_167534/566849597.py", line 21, in captures_leaked_tensor(tf.constant(2)) TypeError: is out of scope and cannot be used here. Use return values, explicit Python locals or TensorFlow collections to access it. Please see https://www.tensorflow.org/guide/function#all_outputs_of_a_tffunction_must_be_return_values for more information. was defined here: File "/usr/lib/python3.9/runpy.py", line 197, in _run_module_as_main File "/usr/lib/python3.9/runpy.py", line 87, in _run_code File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel_launcher.py", line 18, in File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/traitlets/config/application.py", line 1075, in launch_instance File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelapp.py", line 739, in start File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tornado/platform/asyncio.py", line 205, in start File "/usr/lib/python3.9/asyncio/base_events.py", line 601, in run_forever File "/usr/lib/python3.9/asyncio/base_events.py", line 1905, in _run_once File "/usr/lib/python3.9/asyncio/events.py", line 80, in _run File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 545, in dispatch_queue File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 534, in process_one File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 437, in dispatch_shell File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/ipkernel.py", line 362, in execute_request File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 778, in execute_request File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/ipkernel.py", line 449, in do_execute File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/zmqshell.py", line 549, in run_cell File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3048, in run_cell File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3103, in _run_cell File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3308, in run_cell_async File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3490, in run_ast_nodes File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3550, in run_code File "/tmpfs/tmp/ipykernel_167534/566849597.py", line 7, in File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 150, in error_handler File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 833, in __call__ File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 889, in _call File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 696, in _initialize File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 178, in trace_function File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 283, in _maybe_define_function File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 310, in _create_concrete_function File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/func_graph.py", line 1059, in func_graph_from_py_func File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 599, in wrapped_fn File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/autograph_util.py", line 41, in autograph_handler File "/tmpfs/tmp/ipykernel_167534/566849597.py", line 4, in leaky_function File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 150, in error_handler File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/override_binary_operator.py", line 113, in binary_op_wrapper File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/ops/tensor_math_operator_overrides.py", line 28, in _add_dispatch_factory File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 150, in error_handler File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/util/dispatch.py", line 1260, in op_dispatch_handler File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/ops/math_ops.py", line 1701, in _add_dispatch File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/ops/gen_math_ops.py", line 490, in add_v2 File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/op_def_library.py", line 796, in _apply_op_helper File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/func_graph.py", line 670, in _create_op_internal File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 2682, in _create_op_internal File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 1177, in from_node_def The tensor cannot be accessed from here, because it was defined in FuncGraph(name=leaky_function, id=139959630636096), which is out of scope. 2 'SymbolicTensor' object has no attribute 'numpy' Caught expected exception : Traceback (most recent call last): File "/tmpfs/tmp/ipykernel_167534/3551158538.py", line 8, in assert_raises yield File "/tmpfs/tmp/ipykernel_167534/566849597.py", line 21, in captures_leaked_tensor(tf.constant(2)) TypeError: is out of scope and cannot be used here. Use return values, explicit Python locals or TensorFlow collections to access it. Please see https://www.tensorflow.org/guide/function#all_outputs_of_a_tffunction_must_be_return_values for more information. was defined here: File "/usr/lib/python3.9/runpy.py", line 197, in _run_module_as_main File "/usr/lib/python3.9/runpy.py", line 87, in _run_code File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel_launcher.py", line 18, in File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/traitlets/config/application.py", line 1075, in launch_instance File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelapp.py", line 739, in start File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tornado/platform/asyncio.py", line 205, in start File "/usr/lib/python3.9/asyncio/base_events.py", line 601, in run_forever File "/usr/lib/python3.9/asyncio/base_events.py", line 1905, in _run_once File "/usr/lib/python3.9/asyncio/events.py", line 80, in _run File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 545, in dispatch_queue File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 534, in process_one File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 437, in dispatch_shell File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/ipkernel.py", line 362, in execute_request File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/kernelbase.py", line 778, in execute_request File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/ipkernel.py", line 449, in do_execute File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/ipykernel/zmqshell.py", line 549, in run_cell File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3048, in run_cell File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3103, in _run_cell File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/async_helpers.py", line 129, in _pseudo_sync_runner File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3308, in run_cell_async File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3490, in run_ast_nodes File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/IPython/core/interactiveshell.py", line 3550, in run_code File "/tmpfs/tmp/ipykernel_167534/566849597.py", line 7, in File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 150, in error_handler File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 833, in __call__ File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 889, in _call File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 696, in _initialize File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 178, in trace_function File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 283, in _maybe_define_function File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compilation.py", line 310, in _create_concrete_function File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/func_graph.py", line 1059, in func_graph_from_py_func File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/polymorphic_function.py", line 599, in wrapped_fn File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/autograph_util.py", line 41, in autograph_handler File "/tmpfs/tmp/ipykernel_167534/566849597.py", line 4, in leaky_function File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 150, in error_handler File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/override_binary_operator.py", line 113, in binary_op_wrapper File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/ops/tensor_math_operator_overrides.py", line 28, in _add_dispatch_factory File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 150, in error_handler File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/util/dispatch.py", line 1260, in op_dispatch_handler File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/ops/math_ops.py", line 1701, in _add_dispatch File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/ops/gen_math_ops.py", line 490, in add_v2 File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/op_def_library.py", line 796, in _apply_op_helper File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/func_graph.py", line 670, in _create_op_internal File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 2682, in _create_op_internal File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 1177, in from_node_def The tensor cannot be accessed from here, because it was defined in FuncGraph(name=leaky_function, id=139959630636096), which is out of scope. Usually, leaks such as these occur when you use Python statements or data structures. In addition to leaking inaccessible tensors, such statements are also likely wrong because they count as Python side effects, and are not guaranteed to execute at every function call. Common ways to leak local tensors also include mutating an external Python collection, or an object: class MyClass: def __init__(self): self.field = None external_list = [] external_object = MyClass() def leaky_function(): a = tf.constant(1) external_list.append(a) # Bad - leaks tensor external_object.field = a # Bad - leaks tensor class MyClass: def __init__(self): self.field = None external_list = [] external_object = MyClass() def leaky_function(): a = tf.constant(1) external_list.append(a) # Bad - leaks tensor external_object.field = a # Bad - leaks tensor Recursive tf.functions are not supported Recursive tf.functions are not supported and could cause infinite loops. For example, tf.function @tf.function def recursive_fn(n): if n > 0: return recursive_fn(n - 1) else: return 1 with assert_raises(Exception): recursive_fn(tf.constant(5)) # Bad - maximum recursion error. @tf.function def recursive_fn(n): if n > 0: return recursive_fn(n - 1) else: return 1 with assert_raises(Exception): recursive_fn(tf.constant(5)) # Bad - maximum recursion error. Even if a recursive tf.function seems to work, the Python function will be traced multiple times and could have performance implications. For example, tf.function @tf.function def recursive_fn(n): if n > 0: print('tracing') return recursive_fn(n - 1) else: return 1 recursive_fn(5) # Warning - multiple tracings @tf.function def recursive_fn(n): if n > 0: print('tracing') return recursive_fn(n - 1) else: return 1 recursive_fn(5) # Warning - multiple tracings tracing tracing tracing tracing tracing tracing tracing tracing tracing tracing Known Issues If your tf.function is not evaluating correctly, the error may be explained by these known issues which are planned to be fixed in the future. tf.function Depending on Python global and free variables tf.function creates a new ConcreteFunction when called with a new value of a Python argument. However, it does not do that for the Python closure, globals, or nonlocals of that tf.function. If their value changes in between calls to the tf.function, the tf.function will still use the values they had when it was traced. This is different from how regular Python functions work. tf.function ConcreteFunction tf.function tf.function tf.function For that reason, you should follow a functional programming style that uses arguments instead of closing over outer names. @tf.function def buggy_add(): return 1 + foo @tf.function def recommended_add(foo): return 1 + foo foo = 1 print("Buggy:", buggy_add()) print("Correct:", recommended_add(foo)) @tf.function def buggy_add(): return 1 + foo @tf.function def recommended_add(foo): return 1 + foo foo = 1 print("Buggy:", buggy_add()) print("Correct:", recommended_add(foo)) Buggy: tf.Tensor(2, shape=(), dtype=int32) Correct: tf.Tensor(2, shape=(), dtype=int32) Buggy: tf.Tensor(2, shape=(), dtype=int32) Correct: tf.Tensor(2, shape=(), dtype=int32) print("Updating the value of `foo` to 100!") foo = 100 print("Buggy:", buggy_add()) # Did not change! print("Correct:", recommended_add(foo)) print("Updating the value of `foo` to 100!") foo = 100 print("Buggy:", buggy_add()) # Did not change! print("Correct:", recommended_add(foo)) Updating the value of `foo` to 100! Buggy: tf.Tensor(2, shape=(), dtype=int32) Correct: tf.Tensor(101, shape=(), dtype=int32) Updating the value of `foo` to 100! Buggy: tf.Tensor(2, shape=(), dtype=int32) Correct: tf.Tensor(101, shape=(), dtype=int32) Another way to update a global value is to make it a tf.Variable and use the Variable.assign method instead. tf.Variable Variable.assign @tf.function def variable_add(): return 1 + foo foo = tf.Variable(1) print("Variable:", variable_add()) @tf.function def variable_add(): return 1 + foo foo = tf.Variable(1) print("Variable:", variable_add()) Variable: tf.Tensor(2, shape=(), dtype=int32) Variable: tf.Tensor(2, shape=(), dtype=int32) print("Updating the value of `foo` to 100!") foo.assign(100) print("Variable:", variable_add()) print("Updating the value of `foo` to 100!") foo.assign(100) print("Variable:", variable_add()) Updating the value of `foo` to 100! Variable: tf.Tensor(101, shape=(), dtype=int32) Updating the value of `foo` to 100! Variable: tf.Tensor(101, shape=(), dtype=int32) Depending on Python objects Passing custom Python objects as arguments to tf.function is supported but has certain limitations. tf.function For maximum feature coverage, consider transforming the objects into Extension types before passing them to tf.function. You can also use Python primitives and tf.nest-compatible structures. tf.function tf.nest However, as covered in the rules of tracing, when a custom TraceType is not provided by the custom Python class, tf.function is forced to use instance-based equality which means it will not create a new trace when you pass the same object with modified attributes. TraceType tf.function not create a new trace same object with modified attributes class SimpleModel(tf.Module): def __init__(self): # These values are *not* tf.Variables. self.bias = 0. self.weight = 2. @tf.function def evaluate(model, x): return model.weight * x + model.bias simple_model = SimpleModel() x = tf.constant(10.) print(evaluate(simple_model, x)) class SimpleModel(tf.Module): def __init__(self): # These values are *not* tf.Variables. self.bias = 0. self.weight = 2. @tf.function def evaluate(model, x): return model.weight * x + model.bias simple_model = SimpleModel() x = tf.constant(10.) print(evaluate(simple_model, x)) tf.Tensor(20.0, shape=(), dtype=float32) tf.Tensor(20.0, shape=(), dtype=float32) print("Adding bias!") simple_model.bias += 5.0 print(evaluate(simple_model, x)) # Didn't change :( print("Adding bias!") simple_model.bias += 5.0 print(evaluate(simple_model, x)) # Didn't change :( dding bias! tf.Tensor(20.0, shape=(), dtype=float32) dding bias! tf.Tensor(20.0, shape=(), dtype=float32) Using the same tf.function to evaluate the modified instance of the model will be buggy since it still has the same instance-based TraceType as the original model. tf.function For that reason, you're recommended to write your tf.function to avoid depending on mutable object attributes or implement the Tracing Protocol for the objects to inform tf.function about such attributes. tf.function tf.function If that is not possible, one workaround is to make new tf.functions each time you modify your object to force retracing: tf.function def evaluate(model, x): return model.weight * x + model.bias new_model = SimpleModel() evaluate_no_bias = tf.function(evaluate).get_concrete_function(new_model, x) # Don't pass in `new_model`. `tf.function` already captured its state during tracing. print(evaluate_no_bias(x)) def evaluate(model, x): return model.weight * x + model.bias new_model = SimpleModel() evaluate_no_bias = tf.function(evaluate).get_concrete_function(new_model, x) # Don't pass in `new_model`. `tf.function` already captured its state during tracing. print(evaluate_no_bias(x)) tf.Tensor(20.0, shape=(), dtype=float32) tf.Tensor(20.0, shape=(), dtype=float32) print("Adding bias!") new_model.bias += 5.0 # Create new `tf.function` and `ConcreteFunction` since you modified `new_model`. evaluate_with_bias = tf.function(evaluate).get_concrete_function(new_model, x) print(evaluate_with_bias(x)) # Don't pass in `new_model`. print("Adding bias!") new_model.bias += 5.0 # Create new `tf.function` and `ConcreteFunction` since you modified `new_model`. evaluate_with_bias = tf.function(evaluate).get_concrete_function(new_model, x) print(evaluate_with_bias(x)) # Don't pass in `new_model`. Adding bias! tf.Tensor(25.0, shape=(), dtype=float32) Adding bias! tf.Tensor(25.0, shape=(), dtype=float32) As retracing can be expensive, you can use tf.Variables as object attributes, which can be mutated (but not changed, careful!) for a similar effect without needing a retrace. tf.Variable class BetterModel: def __init__(self): self.bias = tf.Variable(0.) self.weight = tf.Variable(2.) @tf.function def evaluate(model, x): return model.weight * x + model.bias better_model = BetterModel() print(evaluate(better_model, x)) class BetterModel: def __init__(self): self.bias = tf.Variable(0.) self.weight = tf.Variable(2.) @tf.function def evaluate(model, x): return model.weight * x + model.bias better_model = BetterModel() print(evaluate(better_model, x)) tf.Tensor(20.0, shape=(), dtype=float32) tf.Tensor(20.0, shape=(), dtype=float32) print("Adding bias!") better_model.bias.assign_add(5.0) # Note: instead of better_model.bias += 5 print(evaluate(better_model, x)) # This works! print("Adding bias!") better_model.bias.assign_add(5.0) # Note: instead of better_model.bias += 5 print(evaluate(better_model, x)) # This works! Adding bias! tf.Tensor(25.0, shape=(), dtype=float32) Adding bias! tf.Tensor(25.0, shape=(), dtype=float32) Creating tf.Variables tf.function only supports singleton tf.Variables created once on the first call, and reused across subsequent function calls. The code snippet below would create a new tf.Variable in every function call, which results in a ValueError exception. tf.function tf.Variable tf.Variable ValueError Example: @tf.function def f(x): v = tf.Variable(1.0) return v with assert_raises(ValueError): f(1.0) @tf.function def f(x): v = tf.Variable(1.0) return v with assert_raises(ValueError): f(1.0) Caught expected exception : Traceback (most recent call last): File "/tmpfs/tmp/ipykernel_167534/3551158538.py", line 8, in assert_raises yield File "/tmpfs/tmp/ipykernel_167534/3018268426.py", line 7, in f(1.0) ValueError: in user code: File "/tmpfs/tmp/ipykernel_167534/3018268426.py", line 3, in f * v = tf.Variable(1.0) ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information. Caught expected exception : Traceback (most recent call last): File "/tmpfs/tmp/ipykernel_167534/3551158538.py", line 8, in assert_raises yield File "/tmpfs/tmp/ipykernel_167534/3018268426.py", line 7, in f(1.0) ValueError: in user code: File "/tmpfs/tmp/ipykernel_167534/3018268426.py", line 3, in f * v = tf.Variable(1.0) ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information. A common pattern used to work around this limitation is to start with a Python None value, then conditionally create the tf.Variable if the value is None: tf.Variable class Count(tf.Module): def __init__(self): self.count = None @tf.function def __call__(self): if self.count is None: self.count = tf.Variable(0) return self.count.assign_add(1) c = Count() print(c()) print(c()) class Count(tf.Module): def __init__(self): self.count = None @tf.function def __call__(self): if self.count is None: self.count = tf.Variable(0) return self.count.assign_add(1) c = Count() print(c()) print(c()) tf.Tensor(1, shape=(), dtype=int32) tf.Tensor(2, shape=(), dtype=int32) tf.Tensor(1, shape=(), dtype=int32) tf.Tensor(2, shape=(), dtype=int32) Using with multiple Keras optimizers You may encounter ValueError: tf.function only supports singleton tf.Variables created on the first call. when using more than one Keras optimizer with a tf.function. This error occurs because optimizers internally create tf.Variables when they apply gradients for the first time. ValueError: tf.function only supports singleton tf.Variables created on the first call. tf.function tf.Variable opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2) opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3) @tf.function def train_step(w, x, y, optimizer): with tf.GradientTape() as tape: L = tf.reduce_sum(tf.square(w*x - y)) gradients = tape.gradient(L, [w]) optimizer.apply_gradients(zip(gradients, [w])) w = tf.Variable(2.) x = tf.constant([-1.]) y = tf.constant([2.]) train_step(w, x, y, opt1) print("Calling `train_step` with different optimizer...") with assert_raises(ValueError): train_step(w, x, y, opt2) opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2) opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3) @tf.function def train_step(w, x, y, optimizer): with tf.GradientTape() as tape: L = tf.reduce_sum(tf.square(w*x - y)) gradients = tape.gradient(L, [w]) optimizer.apply_gradients(zip(gradients, [w])) w = tf.Variable(2.) x = tf.constant([-1.]) y = tf.constant([2.]) train_step(w, x, y, opt1) print("Calling `train_step` with different optimizer...") with assert_raises(ValueError): train_step(w, x, y, opt2) Calling `train_step` with different optimizer... Caught expected exception : Traceback (most recent call last): File "/tmpfs/tmp/ipykernel_167534/3551158538.py", line 8, in assert_raises yield File "/tmpfs/tmp/ipykernel_167534/950644149.py", line 18, in train_step(w, x, y, opt2) ValueError: in user code: File "/tmpfs/tmp/ipykernel_167534/950644149.py", line 9, in train_step * optimizer.apply_gradients(zip(gradients, [w])) File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/optimizers/base_optimizer.py", line 291, in apply_gradients ** self.apply(grads, trainable_variables) File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/optimizers/base_optimizer.py", line 330, in apply self.build(trainable_variables) File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/optimizers/adam.py", line 97, in build self.add_variable_from_reference( File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/backend/tensorflow/optimizer.py", line 36, in add_variable_from_reference return super().add_variable_from_reference( File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/optimizers/base_optimizer.py", line 227, in add_variable_from_reference return self.add_variable( File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/optimizers/base_optimizer.py", line 201, in add_variable variable = backend.Variable( File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/backend/common/variables.py", line 163, in __init__ self._initialize_with_initializer(initializer) File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/backend/tensorflow/core.py", line 40, in _initialize_with_initializer self._value = tf.Variable( ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information. Calling `train_step` with different optimizer... Caught expected exception : Traceback (most recent call last): File "/tmpfs/tmp/ipykernel_167534/3551158538.py", line 8, in assert_raises yield File "/tmpfs/tmp/ipykernel_167534/950644149.py", line 18, in train_step(w, x, y, opt2) ValueError: in user code: File "/tmpfs/tmp/ipykernel_167534/950644149.py", line 9, in train_step * optimizer.apply_gradients(zip(gradients, [w])) File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/optimizers/base_optimizer.py", line 291, in apply_gradients ** self.apply(grads, trainable_variables) File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/optimizers/base_optimizer.py", line 330, in apply self.build(trainable_variables) File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/optimizers/adam.py", line 97, in build self.add_variable_from_reference( File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/backend/tensorflow/optimizer.py", line 36, in add_variable_from_reference return super().add_variable_from_reference( File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/optimizers/base_optimizer.py", line 227, in add_variable_from_reference return self.add_variable( File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/optimizers/base_optimizer.py", line 201, in add_variable variable = backend.Variable( File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/backend/common/variables.py", line 163, in __init__ self._initialize_with_initializer(initializer) File "/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/src/backend/tensorflow/core.py", line 40, in _initialize_with_initializer self._value = tf.Variable( ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information. If you need to change a stateful object between calls, it's simplest to define a tf.Module subclass, and create instances to hold those objects: tf.Module class TrainStep(tf.Module): def __init__(self, optimizer): self.optimizer = optimizer @tf.function def __call__(self, w, x, y): with tf.GradientTape() as tape: L = tf.reduce_sum(tf.square(w*x - y)) gradients = tape.gradient(L, [w]) self.optimizer.apply_gradients(zip(gradients, [w])) opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2) opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3) train_o1 = TrainStep(opt1) train_o2 = TrainStep(opt2) train_o1(w, x, y) train_o2(w, x, y) class TrainStep(tf.Module): def __init__(self, optimizer): self.optimizer = optimizer @tf.function def __call__(self, w, x, y): with tf.GradientTape() as tape: L = tf.reduce_sum(tf.square(w*x - y)) gradients = tape.gradient(L, [w]) self.optimizer.apply_gradients(zip(gradients, [w])) opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2) opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3) train_o1 = TrainStep(opt1) train_o2 = TrainStep(opt2) train_o1(w, x, y) train_o2(w, x, y) You could also do this manually by creating multiple instances of the @tf.function wrapper, one for each optimizer: @tf.function opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2) opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3) # Not a tf.function. def train_step(w, x, y, optimizer): with tf.GradientTape() as tape: L = tf.reduce_sum(tf.square(w*x - y)) gradients = tape.gradient(L, [w]) optimizer.apply_gradients(zip(gradients, [w])) w = tf.Variable(2.) x = tf.constant([-1.]) y = tf.constant([2.]) # Make a new tf.function and ConcreteFunction for each optimizer. train_step_1 = tf.function(train_step) train_step_2 = tf.function(train_step) for i in range(10): if i % 2 == 0: train_step_1(w, x, y, opt1) else: train_step_2(w, x, y, opt2) opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2) opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3) # Not a tf.function. def train_step(w, x, y, optimizer): with tf.GradientTape() as tape: L = tf.reduce_sum(tf.square(w*x - y)) gradients = tape.gradient(L, [w]) optimizer.apply_gradients(zip(gradients, [w])) w = tf.Variable(2.) x = tf.constant([-1.]) y = tf.constant([2.]) # Make a new tf.function and ConcreteFunction for each optimizer. train_step_1 = tf.function(train_step) train_step_2 = tf.function(train_step) for i in range(10): if i % 2 == 0: train_step_1(w, x, y, opt1) else: train_step_2(w, x, y, opt2) Using with multiple Keras models You may also encounter ValueError: tf.function only supports singleton tf.Variables created on the first call. when passing different model instances to the same tf.function. ValueError: tf.function only supports singleton tf.Variables created on the first call. tf.function This error occurs because Keras models (which do not have their input shape defined) and Keras layers create tf.Variables when they are first called. You may be attempting to initialize those variables inside a tf.function, which has already been called. To avoid this error, try calling model.build(input_shape) to initialize all the weights before training the model. do not have their input shape defined tf.Variable tf.function model.build(input_shape) Originally published on the TensorFlow website, this article appears here under a new headline and is licensed under CC BY 4.0. Code samples shared under the Apache 2.0 License. Originally published on the TensorFlow website, this article appears here under a new headline and is licensed under CC BY 4.0. Code samples shared under the Apache 2.0 License. Originally published on the TensorFlow website, TensorFlow website TensorFlow website CC BY 4.0. Apache 2.0 License.