Applying tf.nn.softmax() only to positive elements

Question

I tried far to long to solve this problem and did not find anything useful on the Internet so I have to ask:

Given a tensor T, let's say T = tf.random_normal([100]), I want to apply softmax() only to the positive elements of the tensor. Something like T = tf.nn.softmax(T[T>0]) which of course does not work in Tensorflow.

In short: I want to compute softmax and applied only on elements T > 0.

How can I do that in Tensorflow?

Answer 1

I am new to Tensorflow but this is my try, based on the math formula:

def softmax_positiv(T):
    # softmax = tf.exp(logits) / tf.reduce_sum(tf.exp(logits), axis)
    Tsign = tf.greater(T, 0.)
    Tpos = tf.gather(T, tf.where(Tsign))
    _reduce_sum = tf.reduce_sum(tf.exp(Tpos))
    Tsign = tf.cast(Tsign, tf.float32)    
    Tpos = (tf.exp(T) / _reduce_sum) * Tsign
    Tneg = (Tsign - 1) * -1 * T

  return Tpos+Tneg

---

Updated Version(using @Aldream sugestion):

---

def softmax_positiv(T):
    #softmax = tf.exp(logits) / tf.reduce_sum(tf.exp(logits), axis)
    Tsign=tf.greater(T,0)
    _reduce_sum=tf.reduce_sum(tf.exp(tf.where(Tsign,T,tf.zeros(T.shape))))
    return  tf.where(Tsign, tf.exp(T) / _reduce_sum, T)

Answer 2

If you want softmax computed + applied only to elements T > 0:

An idea could be to create 2 partitions based on your condition (T > 0), apply the operation (softmax) to the target partition, then stitch them back together.

Something like this, using tf.dynamic_partition and tf.dynamic_stitch:

import tensorflow as tf

T = tf.random_normal(shape=(2, 3, 4))

# Creating partition based on condition:
condition_mask = tf.cast(tf.greater(T, 0.), tf.int32)
partitioned_T = tf.dynamic_partition(T, condition_mask, 2)
# Applying the operation to the target partition:
partitioned_T[1] = tf.nn.softmax(partitioned_T[1])

# Stitching back together, flattening T and its indices to make things easier::
condition_indices = tf.dynamic_partition(tf.range(tf.size(T)), tf.reshape(condition_mask, [-1]), 2)
res_T = tf.dynamic_stitch(condition_indices, partitioned_T)
res_T = tf.reshape(res_T, tf.shape(T))

with tf.Session() as sess:
    t, res = sess.run([T, res_T])
    print(t)
    # [[[-1.92647386  0.7442674   1.86053932 -0.95315439]
    #  [-0.38296485  1.19349718 -1.27562618 -0.73016083]
    #  [-0.36333972 -0.90614134 -0.15798278 -0.38928652]]
    # 
    # [[-0.42384467  0.69428021  1.94177043 -0.13672788]
    #  [-0.53473723  0.94478583 -0.52320045  0.36250541]
    #  [ 0.59011376 -0.77091616 -0.12464728  1.49722672]]]
    print(res)
    # [[[-1.92647386  0.06771058  0.20675084 -0.95315439]
    #  [-0.38296485  0.10610957 -1.27562618 -0.73016083]
    #  [-0.36333972 -0.90614134 -0.15798278 -0.38928652]]
    # 
    # [[-0.42384467  0.06440912  0.22424641 -0.13672788]
    #  [-0.53473723  0.08274478 -0.52320045  0.04622314]
    #  [ 0.05803747 -0.77091616 -0.12464728  0.14376813]]]

---

Previous answer

This answer is valid only if you want softmax to be computed over all elements of T but applied only to those greater than 0.

Using tf.where():

T = tf.where(tf.greater(T, 0.), tf.nn.softmax(T), T)