Computing inverse of a Tensor using tf.map_fn

Question

Using Tensorflow 1.4 I want to compute elementwise inverse (x --> 1/x) of a tensor using map functions. If the value of an element in the Tensor is zero, I want to have zero output.

As an example for the tensor: [[0, 1, 0], [0.5, 0.5, 0.3]], I want to have the output: [[0,1,0], [2, 2, 3.333]]. I know that I can easily get the desired output using tf.math.reciprocal_no_nan() in tf2.0 and tf.math.divide_no_nan() in tf 1.4, but I am wondering why the following code does not work:

tensor = tf.constant([[0, 1, 0], [0.5, 0.5, 0.3]], tf.float32)
tensor_inverse = tf.map_fn(lambda x: tf.cond(tf.math.not_equal(x, 0.0), lambda x: 1/x, lambda: 0) , tensor)

I get this error:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Answer 1

Let us break down your code example:

First function you are using is map_fn. Map_fn will split the tensor across the first dimension and pass those individual tensors to it's internal supplied function. This function is not causing any problem for you. Next is tf.cond. tf.cond expects a scalar value in it's predicate. To break down:

tensor = tf.constant([[0, 1, 0], [0.5, 0.5, 0.3]], tf.float32)
cond_val1 = tf.math.not_equal(tensor, 0.0)
print(cond_val1.shape)  # Shape (2, 3)

cond_val1 in above example is clearly a tensor. You will have to make use of tf.reduce_all or tf.reduce_any to convert it into a scalar. Then you will be getting a scalar as required for tf.cond. For ex:

cond_val2 = tf.reduce_all(tf.math.not_equal(tensor, 0.0))
print(cond_val2.shape)  # Shape ()

Now this will make tf.cond work. But there is another problem you are getting. You have already lost your ability to process the tensor element-wise.

Secondly, by means of map_fn you are passing the entire tensor split at first dimension which in your case will be [0, 1, 0] and [0.5, 0.5, 0.3]. But the problem is your tf.cond's true_fn and false_fn don't have ability to process it element-wise.

Hope you got to know the variety of problems which are present in your code.