Handle size-0 output for split (facebookincubator#1011)

Summary:
Pull Request resolved: facebookincubator#1011

For size-0 output tensors, we do not allocate space for the tensor, therefore a nullptr will be passed in.

Reviewed By: frank-wei

Differential Revision: D59128217

fbshipit-source-id: f1c761f37bb91049d9ac6726464c4a4c7881889c

python/aitemplate/backend/common/split_common.py

@@ -386,7 +386,7 @@
     throw std::runtime_error("input is NULL!");
   }
   for (int i = 0; i < real_num_splits; i++) {
-    if (!outputs[i]) {
+    if (split_sizes[i] && !outputs[i]) {
       throw std::runtime_error("NULL output found at: " + std::to_string(i));
     }
   }

tests/unittest/ops/test_split.py

@@ -204,6 +204,7 @@ def test_split(self):
         self._run_split(input_shape=[2, 0, 4], split_size_or_sections=0, dim=-2)
         self._run_split(input_shape=[2, 0, 4], split_size_or_sections=2, dim=-1)
         self._run_split(input_shape=[2, 0, 7], split_size_or_sections=[2, 3, 2], dim=-1)
+        self._run_split(input_shape=[32, 8], split_size_or_sections=[8, 0, 0], dim=-1)
 
     def test_split_with_mask(self):
         self._run_split(