title: Applying softmax to non-zero elements in the matrix across a dimension
body: <p>Perhaps this is trivial, but perhaps it is not. I have spent way too much time trying to figure out how to make this work. Here is the code:</p>
  
  <pre><code># batch x time x events
  batch = 2
  time = 3
  events = 4
  tensor = np.random.rand(batch, time, events)
  
  tensor[0][0][2] = 0
  tensor[0][0][3] = 0
  
  tensor[0][1][3] = 0
  
  tensor[0][2][1] = 0
  tensor[0][2][2] = 0
  tensor[0][2][3] = 0
  
  tensor[1][0][3] = 0
  
  non_zero = ~tf.equal(tensor, 0.)
  
  s = tf.Session()
  g = tf.global_variables_initializer()
  s.run(g)
  
  s.run(non_zero)
  </code></pre>
  
  <p>I am trying to apply <code>tf.nn.softmax</code> to the non-zero values across each of the <code>time</code> dimensions. However, when I am using <code>tf.boolean_mask</code> then it actually gathers all of the non-zero values together. That is not what I want. I want to preserve the dimensions.</p>
  
  <p>Here is the screenshot of what the tensor looks like:
  <a href="https://i.stack.imgur.com/4I6bC.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/4I6bC.png" alt="enter image description here"></a></p>
  
  <p>So <code>tf.nn.softmax</code> should be applied to only those groups and it should "put them back" into their original positions. Does anyone know how to do this?</p>
  
  <p><strong>EDIT:</strong></p>
  
  <p>I almost found a solution that I need, with your help guys. But I am still missing one step. Assigning the softmax across each time dimension to the non-zero values:</p>
  
  <pre><code>def apply_sparse_softmax(time_vector):
      non_zeros = ~tf.equal(time_vector, 0.)
  
      sparse_softmax = tf.nn.softmax(tf.boolean_mask(time_vector, non_zeros))
      new_time_vector = sparse_softmax * tf.cast(non_zeros, tf.float64) # won't work because dimensions are different
      return time_vector
  </code></pre>
  
  <p>Please also note that this solution should handle the cases when you have zeros all across the time dimension. Then it should just stay the same.</p>