jknielse
jknielse

Reputation: 751

How do I zip tensors in tensorflow when the dimensions don't match

I have two tensors, one of shape [None, 20, 2], and one of shape [None, 1]. I would like to do an operation on each of the sub-tensors in lockstep to produce a value such that I would end up with a tensor of shape [None, 1].

In python land, I would zip these two, and iterate over the result.

So, just to be clear, I'd like to write a function that takes a [20, 2]-shape tensor and a [1]-shape tensor, and produces a [1]-shape tensor, then apply this function to the [None, 20, 2] and [None, 1] tensors, to produce a [None, 1] tensor.

Hope I articulated that well enough; higher dimensionality makes my head spin sometimes.

Upvotes: 1

Views: 1549

Answers (1)

bask0
bask0

Reputation: 360

This works for me (TensorFlow version 1.4.0)

tf.reset_default_graph()
sess = tf.Session()

# Define placeholders with undefined first dimension.
a = tf.placeholder(dtype=tf.float32, shape=[None, 3, 4])
b = tf.placeholder(dtype=tf.float32, shape=[None, 1])

# Create some input data.
a_input = np.arange(24).reshape(2, 3, 4)
b_input = np.arange(2).reshape(2, 1)

# TensorFlow map function.
def f_tf(x):
    return tf.reduce_sum(x[0]) + tf.reduce_sum(x[1])
# Numpy map function (for validation of results).
def f_numpy(x):
    return np.sum(x[0]) + np.sum(x[1])

# Run TensorFlow function.
s = tf.map_fn(f, [a, b], dtype=tf.float32)
sess.run(s, feed_dict={a: a_input, b: b_input})

array([ 66., 211.], dtype=float32)

# Run Numpy function. 
for inp in zip(a_input, b_input):
    print(f_numpy(inp))

66 211

Upvotes: 2

Related Questions