I don't understand map_fn with two inputs

Question

I'm having trouble understanding what map_fn does, when I use it on a tuple. For testing I did the following thing:

a = tf.constant([[1,2,3,4],[5,6,7,8],[9,10,11,12]])
b = tf.constant([[1,2,3],[4,5,6],[7,8,9]])
func = lambda x: x[0]
y = tf.map_fn(func,(a,b))

What I expected here is that map_fn takes a and b apart into three vectors each and gives them to func. func then only returns the first input and map_fn stacks them back together. So I thought I should just get a back. What happens instead is a terrible error:

ValueError: The two structures don't have the same nested structure.

First structure: type=tuple str=(<tf.Tensor: shape=(3, 4), dtype=int32, numpy=
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]], dtype=int32)>, <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]], dtype=int32)>)

Second structure: type=EagerTensor str=tf.Tensor([1 2 3 4], shape=(4,), dtype=int32)

More specifically: Substructure "type=tuple str=(<tf.Tensor: shape=(3, 4), dtype=int32, numpy=
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]], dtype=int32)>, <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]], dtype=int32)>)" is a sequence, while substructure "type=EagerTensor str=tf.Tensor([1 2 3 4], shape=(4,), dtype=int32)" is not

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
9 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py in assert_same_structure(nest1, nest2, check_types, expand_composites)
    400                   "Entire first structure:\n%s\n"
    401                   "Entire second structure:\n%s"
--> 402                   % (str(e), str1, str2))
    403 
    404 

ValueError: The two structures don't have the same nested structure.

First structure: type=tuple str=(<tf.Tensor: shape=(3, 4), dtype=int32, numpy=
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]], dtype=int32)>, <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]], dtype=int32)>)

Second structure: type=EagerTensor str=tf.Tensor([1 2 3 4], shape=(4,), dtype=int32)

More specifically: Substructure "type=tuple str=(<tf.Tensor: shape=(3, 4), dtype=int32, numpy=
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]], dtype=int32)>, <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]], dtype=int32)>)" is a sequence, while substructure "type=EagerTensor str=tf.Tensor([1 2 3 4], shape=(4,), dtype=int32)" is not
Entire first structure:
(., .)
Entire second structure:
.

It seems to me that map_fn somehow tries to combine the whole tuple with one of it's components. Can someone please explain what's going on here?

Answer 1

From the docs:

If fn's input and output signatures are different, then the output signature must be specified using fn_output_signature.

Your function takes a tuple of two tensors and returns a single tensor, so you need to specify fn_output_signature:

y = tf.map_fn(func,(a,b), fn_output_signature=tf.int32)