[TF:XLA] Add support for negative values of "split_dim" argument to Split operator.

PiperOrigin-RevId: 170755169

2 files changed, 26 insertions, 23 deletions

diff --git a/tensorflow/compiler/tests/binary_ops_test.py b/tensorflow/compiler/tests/binary_ops_test.py
index f3ea57596e..792c01327c 100644
--- a/tensorflow/compiler/tests/binary_ops_test.py
+++ b/tensorflow/compiler/tests/binary_ops_test.py
@@ -790,28 +790,30 @@ class BinaryOpsTest(XLATestCase):
 
   def testSplit(self):
     for dtype in self.numeric_types:
-      self._testBinary(
-          lambda x, y: array_ops.split(value=y, num_or_size_splits=3, axis=x),
-          np.int32(0),
-          np.array([[[1], [2]], [[3], [4]], [[5], [6]]],
-                   dtype=dtype),
-          expected=[
-              np.array([[[1], [2]]], dtype=dtype),
-              np.array([[[3], [4]]], dtype=dtype),
-              np.array([[[5], [6]]], dtype=dtype),
-          ],
-          equality_test=self.ListsAreClose)
-
-      self._testBinary(
-          lambda x, y: array_ops.split(value=y, num_or_size_splits=2, axis=x),
-          np.int32(1),
-          np.array([[[1], [2]], [[3], [4]], [[5], [6]]],
-                   dtype=dtype),
-          expected=[
-              np.array([[[1]], [[3]], [[5]]], dtype=dtype),
-              np.array([[[2]], [[4]], [[6]]], dtype=dtype),
-          ],
-          equality_test=self.ListsAreClose)
+      for axis in [0, -3]:
+        self._testBinary(
+            lambda x, y: array_ops.split(value=y, num_or_size_splits=3, axis=x),
+            np.int32(axis),
+            np.array([[[1], [2]], [[3], [4]], [[5], [6]]],
+                     dtype=dtype),
+            expected=[
+                np.array([[[1], [2]]], dtype=dtype),
+                np.array([[[3], [4]]], dtype=dtype),
+                np.array([[[5], [6]]], dtype=dtype),
+            ],
+            equality_test=self.ListsAreClose)
+
+      for axis in [1, -2]:
+        self._testBinary(
+            lambda x, y: array_ops.split(value=y, num_or_size_splits=2, axis=x),
+            np.int32(axis),
+            np.array([[[1], [2]], [[3], [4]], [[5], [6]]],
+                     dtype=dtype),
+            expected=[
+                np.array([[[1]], [[3]], [[5]]], dtype=dtype),
+                np.array([[[2]], [[4]], [[6]]], dtype=dtype),
+            ],
+            equality_test=self.ListsAreClose)
 
   def testTile(self):
     for dtype in self.numeric_types:
diff --git a/tensorflow/compiler/tests/randomized_tests.cc b/tensorflow/compiler/tests/randomized_tests.cc
index b3ec9424c7..7e307f16af 100644
--- a/tensorflow/compiler/tests/randomized_tests.cc
+++ b/tensorflow/compiler/tests/randomized_tests.cc
@@ -2653,7 +2653,8 @@ TEST_F(OpTest, Split) {
     std::vector<int64> dims = RandomDims(1);
     std::uniform_int_distribution<int> ud;
     int32 dim = std::uniform_int_distribution<int32>(
-        0, static_cast<int32>(dims.size()) - 1)(generator());
+        -static_cast<int32>(dims.size()),
+        static_cast<int32>(dims.size()) - 1)(generator());
     int n = std::uniform_int_distribution<int>(1, 5)(generator());
     // Ensure 'dim' is evenly divisible by 'n'.
     dims[dim] /= n;