tf.py_func() unexpected outputs in loop using lambda function

Question

I am observing diverging behaviour from both Numpy and pure Tensorflow when implementing the same simple functionality that shares a variable across iterations in a for loop with tf.py_func.

Let's start with the pure Numpy version:

def my_func(x, k):
    return np.tile(x,k)

x = np.ones((1), np.int64)
for i in range(1,3):
    x = my_func(x, i)

print(x)

This produces the expected output. Initially x is [1]. On the first iteration, it is replicated once to produce [1]. Then on the next iteration, the result is replicated twice, producing the final output [1 1].

A similar approach also produces the same expected output in pure Tensorflow:

x = tf.constant([1], tf.int64)
for i in range(1,3):
    x = tf.tile(x, [i])

with tf.Session() as sess:
    xx = sess.run(x)
    print(xx)

The output is [1 1].

Now I am attempting to do essentially the same thing using tf.py_func, and I can't wrap my head around why I am seeing a different output. This code:

import tensorflow as tf
import numpy as np

def my_func(x, k):
    return np.tile(x,k)

x = tf.constant([1], tf.int64)
for i in range(1,3):
    x = tf.py_func(lambda y: my_func(y, i), [x], tf.int64)

with tf.Session() as sess:
    xx = sess.run(x)
    print(xx)

Produces the unexpected result [1 1 1 1].

Why is this happening? Does py_func have some property that it does not work well with sharing (tensor) variable names, in this case the variable x that is updated on each loop iteration?

Please note that this is a simplified example that reproduces the problem, and whose functionality is easy to reproduce in pure Tensorflow. In my actual application there is a need to use tf.py_func, since the functionality is more complicated.

BlueSun · Accepted Answer · 2018-04-17 11:19:18Z

2

Without the lambda function, it works as intended:

import tensorflow as tf
import numpy as np

def my_func(x, k):
    return np.tile(x,k)

x = tf.constant([1], tf.int64)
for i in range(1,3):
    x = tf.py_func(my_func, [x, i], tf.int64)

with tf.Session() as sess:
    xx = sess.run(x)
    print(xx)

returns [1 1]

edit

I found out why: lambda y: my_func(y, i) saves i by reference and not by value. Therefore the last i value of the for loop is applied to all py_func in the loop. Here is a simpler example that shows the problem:

import tensorflow as tf

def my_func(x, y):
  return x - y

x1 = tf.constant([0], tf.float32)
for i in range(2):
    x1 = tf.py_func(lambda y: my_func(y, i), [x1], tf.float32)

x2 = tf.constant([0], tf.float32)
x2 = tf.py_func(lambda y: my_func(y, 0), [x2], tf.float32)
x2 = tf.py_func(lambda y: my_func(y, 1), [x2], tf.float32)

with tf.Session() as sess:
    print(sess.run(x1))
    print(sess.run(x2))

edited Apr 17, 2018 at 11:19

answered Mar 29, 2018 at 12:41

BlueSun

3,5701 gold badge21 silver badges37 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

mikkola Over a year ago

Nice job! This one really had me scratching my head.

Collectives™ on Stack Overflow

tf.py_func() unexpected outputs in loop using lambda function

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related