Fixes platform dependency in soft_round_ops unit test.

The gradient value of the inverse soft round function at 0 is allowed to
evaluate as infinite, but can sometimes be finite depending on the platform.
This change makes that point a don't-care value.

PiperOrigin-RevId: 361500720
Change-Id: I0cf42414d41f67b66b5fd25f5e723a2f23324c64

Author: Johannes Ballé

tensorflow_compression/python/ops/soft_round_ops_test.py

@@ -77,13 +77,14 @@ def test_soft_round_inverse_values_and_gradients_are_finite(self, alpha):
       y = soft_round_ops.soft_round_inverse(x, alpha=alpha)
     dy = tape.gradient(y, x)
     self.assertAllEqual(tf.math.is_finite(y), tf.ones(x.shape, dtype=bool))
+    is_finite = tf.math.is_finite(dy)
+    expected_finite = tf.ones(dy.shape, dtype=bool)
     if alpha > 15:
-      # We allow non-finite values for large alphas, since the function simply
-      # is extremely steep there.
-      expected_finite = tf.one_hot(5, 11, False, True)
-    else:
-      expected_finite = tf.ones(x.shape, dtype=bool)
-    self.assertAllEqual(tf.math.is_finite(dy), expected_finite)
+      # We allow non-finite values at 0 for large alphas, since the function
+      # simply is extremely steep there.
+      expected_finite = tf.tensor_scatter_nd_update(
+          expected_finite, [[5]], [is_finite[5]])
+    self.assertAllEqual(is_finite, expected_finite)
 
 
 if __name__ == "__main__":