title: Apply a function (tf.square()) to some values of a Tensor - TensorFlow
body: <p>Let's start with an example. I have a tensor with shape of <code>[3, 3, 3]</code> like this</p>
  
  <pre><code>input = tf.constant([[[1.2, -2.3, 4.5], [1.03, 2.89, -1.2], [2.1, 1.9, -1.5]],
                       [[1.3, -2.4, 4.6], [1.04, 2.88, -1.3], [2.2, 1.8, -1.6]],
                       [[1.4, -2.5, 4.7], [1.05, 2.87, -1.4], [2.3, 1.7, -1.7]]],
                      dtype=tf.float32)
  </code></pre>
  
  <p>and I want to apply <code>tf.square</code> only to the elements lower than <code>2</code> for example.
  What I'm doing is this</p>
  
  <pre><code>indices = tf.where(input &lt;= 2)
  base = tf.zeros_like(input)
  ones = tf.constant(1, shape=[18])
  mask = tf.scatter_nd(indices, ones, tf.cast(tf.shape(base), indices.dtype))
  mask = tf.cast(mask, tf.float32)
  masked = tf.square(tf.multiply(input, mask))
  neg_mask = 1 - mask
  neg_masked = tf.multiply(input, neg_mask)
  output = tf.add(masked, neg_masked)
  </code></pre>
  
  <p>and it works. The final output is this</p>
  
  <pre><code>output = tf.constant([[[1.44, 5.29, 4.5], [1.0609, 2.89, 1.44], [2.1, 3.61, 2.25]],
                        [[1.69, 5.76, 4.6], [1.0816, 2.88, 1.69], [2.2, 3.24, 2.56]],
                        [[1.96, 6.25, 4.7], [1.1025, 2.87, 1.96], [2.3, 2.89, 2.89]]],
                       dtype=tf.float32)
  </code></pre>
  
  <p>The problem is that it's quite tricky because this is a toy example but in my case, the tensor has a shape of thousands of elements. Moreover as you can see in this line <code>ones = tf.constant(1, shape=[18])</code> I have coded <code>18</code> because I know that they are 18 and if I try to use for example <code>ones = tf.constant(1, shape=[indices.get_shape()[0]])</code> I get this error <code>TypeError: <strong>long</strong> returned non-long (type NoneType)</code>.</p>
  
  <p>So my two questions are:</p>
  
  <ol>
  <li>How can I solve this problem of the shape?</li>
  <li>Is there a faster and more efficient way to do this?</li>
  </ol>
  
  <p>Thanks in advance</p>
  
  <hr>
  
  <p><strong>edit</strong></p>
  
  <p>Question 1 solved doing this <code>ones = tf.ones(shape=[tf.shape(indices)[0]])</code></p>