Fix tensordot with list of ints as axes (#11959)

* Fix tensordot with list of ints as axes

This fix tries to fix the tensordot issue raised in 11950 where
```
TypeError: object of type 'int' has no len()
```
is thrown out when the axes is a list of ints:
```
with tf.Session() as sess:
    sess.run(
         tf.tensordot(tf.ones((3,3), dtype=tf.float32), tf.constant([2, 3, 1], dtype=tf.float32)[None, None], axes=[1, 2])
    )
```

In numpy it is possible to specify `np.tensordot(np.ones((3,3)), np.array([2, 3, 1])[None, None], axes=[1, 2])`

This fix fixes #11950.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Add test cases for tensordot with axes of list of ints

This commit adds test cases for tensordot with axes of list of ints

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

1 files changed, 17 insertions, 0 deletions

diff --git a/tensorflow/python/kernel_tests/tensordot_op_test.py b/tensorflow/python/kernel_tests/tensordot_op_test.py
index 71230ba000..f375157287 100644
--- a/tensorflow/python/kernel_tests/tensordot_op_test.py
+++ b/tensorflow/python/kernel_tests/tensordot_op_test.py
@@ -20,6 +20,7 @@ from __future__ import print_function
 
 import numpy as np
 
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors_impl
 from tensorflow.python.ops import array_ops
@@ -84,6 +85,22 @@ class TensordotTest(test_lib.TestCase):
                                    b_ph: b,
                                    axes_ph: axes_value})
 
+  # Test case for 11950
+  def test_valid_axis(self):
+    for axes_value in [1, 2], [[1], [2]]:
+      with self.test_session() as sess:
+        np_a = np.ones((3,3))
+        np_b = np.array([2, 3, 1])[None, None]
+        np_ans = np.tensordot(np_a, np_b, axes_value)
+
+        tf_a = array_ops.ones((3,3), dtype=dtypes.float32)
+        tf_b = constant_op.constant([2, 3, 1], dtype=dtypes.float32)[None, None]
+        tf_ans = math_ops.tensordot(tf_a, tf_b, axes_value).eval()
+
+        self.assertAllEqual(tf_ans.shape, np_ans.shape)
+        self.assertAllEqual(tf_ans, np_ans)
+
+
   def test_partial_shape_inference(self):
     a = array_ops.placeholder(dtypes.float32)
     b = array_ops.placeholder(dtypes.float32)