Xiang Zhang
Xiang Zhang

Reputation: 2973

how to customize an element-wise function on tensor in tensorflow?

Say, I have a tensor, it might contain positive and negative values:

[ 1, -1, 2, -2 ]

Now, I want to apply log(x) for positive values, and a constant -10 for negative values:

[ log(1), -10, log(2), -10 ]

In another word, I want to have a function like numpy.vectorize. Is this possible in tensorflow?

One possible way is to use a non-learnable variable, but I don't know if it can properly do back propagation.

Upvotes: 1

Views: 2148

Answers (2)

mrry
mrry

Reputation: 126184

tf.map_fn() enables you to map an arbitrary TensorFlow subcomputation across the elements of a vector (or the slices of a higher-dimensional tensor). For example:

a = tf.constant([1.0, -1.0, 2.0, -2.0])

def f(elem):
  return tf.where(elem > 0, tf.log(elem), -10.0)

  # Alternatively, if the computation is more expensive than `tf.log()`, use
  # `tf.cond()` to ensure that only one branch is executed:
  # return tf.where(elem > 0, lambda: tf.log(elem), lambda: -10.0)

result = tf.map_fn(f, a)

Upvotes: 3

Xiang Zhang
Xiang Zhang

Reputation: 2973

I found it, the tf.where does exactly this kind of job: https://www.tensorflow.org/api_docs/python/tf/where

Upvotes: 0

Related Questions