make tf.nn.softmax axis=-2 safer

mingkaic · Sep 26, 2019 · 25d3e41 · 25d3e41
1 parent 84633d7
commit 25d3e41
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/eteq/test/ptest.py b/eteq/test/ptest.py
@@ -427,7 +427,7 @@ def test_softmax(self):
             self._common_unary_tf(shape, lambda arr: tc.softmax(arr,
                 offset=0, ndims=1), tf.nn.softmax)
             self._common_unary_tf(shape, lambda arr: tc.softmax(arr,
-                offset=1, ndims=1), lambda arr: tf.nn.softmax(arr, axis=-2))
+                offset=1, ndims=1), lambda arr: tf.nn.softmax(arr, axis=len(shape)-2))
 
     def test_square(self):
         shapes = [[3, 4, 5]]