zip like function in Tensorflow? Tensorflow tensor operation

Question

My question is about the tensor operation in Tensorflow. Let's say:

import tensorflow as tf
import numpy as np

a = tf.Variable(np.random.random([10, 3, 3]))
b = tf.Variable(np.random.random([10, 3, 3]))

def some_function(m,n):
    # just as an example
    return tf.add(m, n)

This works in Tensorflow but it requires to know the dimension in advanced. However, it is very likely that the first dimension of the Tensor is None.

c = []
for i in range(10):
    c.append(some_function(a[i], b[i]))
c = tf.stack(c)

So I wonder if there is a zip-like function in Tensorflow? Then we can do:

# TypeError: zip argument #1 must support iteration
c = []
for i, j in zip(a,b):
    c.append(some_function(i,j))
c = tf.stack(c)

Maybe we can use some function like tf.map_fn or tf.scan? But I am not sure. Really thank you, guys.

Answer 1

Tensor objects are not iterable, which explains why your third code sample fails. So, to answer your question, there is no zip-like function in TensorFlow.

You can indeed use tf.map_fn to apply a function to a sequence of tensors. The problem you pose in your example code can be solved in the following fashion:

def some_function(tensor):
  return tf.reduce_sum(tensor)

c = tf.stack([a, b], axis=1)
d = tf.map_fn(some_function, c, dtype=tf.float32)

yields a Tensor d whose value is [20., 6., 6.].