title: tensorflow gradient - getting all nan values
body: <p>I am using python 3 with anaconda, and tensorflow 1.12 with eager eval.</p>
  
  <p>I am using it to create a triplet loss function for a siamese network, and need to calculate distance between different data samples.</p>
  
  <p>I created a function in order to create the distance calculation, but no matter what I do, when I try to calculate it's gradient with respect to the networks output, It keeps giving me all nan gradient.</p>
  
  <p>This is the code:</p>
  
  <pre><code>def matrix_row_wise_norm(matrix):
      import tensorflow as tf
      tensor = tf.expand_dims(matrix, -1)
  
      tensor = tf.transpose(tensor, [0, 2, 1]) - tf.transpose(tensor, [2, 0, 1])
      norm = tf.norm(tensor, axis=2)
      return norm
  </code></pre>
  
  <p>In the loss function I am using </p>
  
  <pre><code>def loss(y_true, p_pred):
      with tf.GradientTape() as t:
      t.watch(y_pred)
          distance_matrix = matrix_row_wise_norm(y_pred)
          grad = t.gradient(distance_matrix, y_pred)
  </code></pre>
  
  <p>And the grad is all <code>nan</code>s. 
  I checked that <code>y_pred</code> is made of legit values - and it does. 
  I tried to create a gradient of <code>y_pred * 2</code> with respect to itself and got legitimate gradient values. </p>
  
  <p>What am I missing here? Is the indexing in the creation of the distance matrix problematic?</p>
  
  <hr>
  
  <p>edit:</p>
  
  <p>the dtype of both <code>y_pred</code> and <code>loss</code> is <code>tf.float32</code></p>
  
  <p>edit: found an <a href="https://github.com/tensorflow/tensorflow/issues/12071" rel="noreferrer">open bug report</a> in tf - could this be the issue? </p>
  
  <hr>
  
  <p>edit:</p>
  
  <p>When I change the norm axis to 0 or 1, I am getting legitimate values and nothing goes to <code>nan</code>. The operation I am getting using norm with <code>axis=2</code> is the pairwise distance between the pairs of rows in the matrix, I suspected this might have something to do with 0 distance between a row to itself, so I clipped the values with min value of 1e-7 without any luck.</p>
  
  <p>Thanks</p>