Tensorflow: Apply different function mapping to different axis of a tf array?

user_12238214 0
Viewed 104

I am new to tensorflow 1.14. Say I have an n * dim tf array: X_tf. I also have dim different numbers of f (as each of them imposes different a and b):

def f(x, a, b):
    return x + a + b

I would like to do the following mapping:

result = [f_0(X_tf[:, 0], a_0, b_0), f_1(X_tf[:, 1], a_1, b_1), ..., f_n(X_tf[:, n], a_n, b_n)]

I see tf.map_fn is capable of doing similar things but the difference is I need to apply for every column of X_tf. While I am also seeking a way to avoid using tf.while_loop and I have to keep these within tf framework.

Any hint on this?

tensorflow python
1 Answers

Make use of tf.stack to split your tensor in first axis and then make use of tf.unstack to join them together. Here is a code sample to do your task.

import tensorflow as tf
import numpy as np

tf.enable_eager_execution()

def f0(X, a, b):
    return X - a - b

def f1(X, a, b):
    return X - a + b

def f2(X, a, b):
    return X + a - b

def f3(X, a, b):
    return X + a + b

X_tf = tf.convert_to_tensor(np.random.randn(4, 10, 10))
a = tf.convert_to_tensor(np.random.randn(4, 10))
b = tf.convert_to_tensor(np.random.randn(4, 10))

X_tf_slices = tf.unstack(X_tf, axis=0)
print(len(X_tf_slices), X_tf_slices[0].shape)   # 4, (10, 10)

a_slices = tf.unstack(a, axis=0)
b_slices = tf.unstack(b, axis=0)

fn_list = [f0, f1, f2, f3]

X_tf = tf.stack([f_i(X_tf_i, a_i, b_i) for f_i, X_tf_i, a_i, b_i in
                 zip(fn_list, X_tf_slices, a_slices, b_slices)])

print(X_tf.shape)   # (4, 10, 10)
user_11908820
Related