CODEWITHSUNDEEP
tensorflow
David Parks
David Parks

Reputation: 32051

Can I apply tf.map_fn(...) to multiple inputs/outputs?

a = tf.constant([[1,2,3],[4,5,6]])
b = tf.constant([True, False], dtype=tf.bool)

a.eval()
array([[1, 2, 3],
       [4, 5, 6]], dtype=int32)
b.eval()
array([ True, False], dtype=bool)

I want to apply a functions to the inputs above, a, and b using tf.map_fn. It will input both [1,2,3], and True and output similar values.

Let's say out function is simply the identity: lambda(x,y): x,y so, given an input of [1,2,3], True, it will output those identical tensors.

I know how to use tf.map_fn(...) with one variable, but not with two. And in this case I have mixed data types (int32 and bool) so I can't simply concatenate the tensors and split them after the call.

Can I use tf.map_fn(...) with multiple inputs/outputs of different data types?

Upvotes: 17

Views: 15413

Answers (4)

Louis Yang
Louis Yang

Reputation: 3821

Updated answer for TFv2

Given

a = tf.constant([[1,2,3],[4,5,6]], dtype=tf.int32)
b = tf.constant([True, False], dtype=tf.bool)

The fn should only take 1 arg but you can pass a tuple of tf.Tensor to your function and then split them inside the function like

def fn(a_and_b):
    a, b = a_and_b
    # then do more stuff on a and b
    return (a, b)

Then call tf.map_fn by

c = tf.map_fn(
    fn,
    elems=(a, b),
    fn_output_signature=(
        tf.TensorSpec(shape=[3], dtype=tf.int32), 
        tf.TensorSpec(shape=[], dtype=tf.bool), 
    ),
)

This gives

>>> c
(<tf.Tensor: shape=(2, 3), dtype=int32, numpy=
 array([[1, 2, 3],
        [4, 5, 6]], dtype=int32)>,
 <tf.Tensor: shape=(2,), dtype=bool, numpy=array([ True, False])>)

Note dtype argument is deprecated for tf.map_fn.

Upvotes: 1

Dylan Madisetti
Dylan Madisetti

Reputation: 775

Both of these are a little messy, you can define the stack "type" explicitly with a custom Extension Type. Here's your example in a hopefully more readable form:

class BoolAndVector(tf.experimental.BatchableExtensionType):
    """A collection bools and vectors."""
    bool: tf.Tensor
    vector: tf.Tensor

@tf.function
def fn(bool_and_vector):
   return BoolAndVector(bool=bool_and_vector.bool,
                         vector=[bool_and_vector.vector[0]])

c = tf.map_fn(fn, BoolAndVector(bool=b, vector=a))

You might have to explicitly state the signature type, but the errors will guide you to what TF wants.

Note, you'll need to use BatchableExtensionType to make this usable with map.

Upvotes: 0

Rocketlife2810
Rocketlife2810

Reputation: 19

Do what is said below, but tuple unpack the values inside the tf.function without the lambda, as computationally expensive and lambdas do not work well with TensorFlow as tf functions.

@tf.function
def x(x):
  x,y = tuple(x) 
  return 

c = tf.map_fn(, (a,b), dtype=(tf.int32, tf.bool))

Upvotes: -1

David Parks
David Parks

Reputation: 32051

Figured it out. You have to define the data types for each tensor in dtype for each of the different tensors, then you can pass the tensors as a tuple, your map function receives a tuple of inputs, and map_fn returns back back a tuple.

Example that works:

a = tf.constant([[1,2,3],[4,5,6]])
b = tf.constant([True, False], dtype=tf.bool)

c = tf.map_fn(lambda x: (x[0], x[1]), (a,b), dtype=(tf.int32, tf.bool))

c[0].eval()
array([[1, 2, 3],
       [4, 5, 6]], dtype=int32)
c[1].eval()
array([ True, False], dtype=bool)

Upvotes: 24

Related Questions