From efd35384db04356d511b5f6fae50f3dd091ea224 Mon Sep 17 00:00:00 2001
From: Jiawei Wang <wangjiawei04@baidu.com>
Date: Mon, 27 Sep 2021 20:40:01 +0800
Subject: [PATCH] fix zero tensor for unique, unstack (#36021)

* fix extra op for expand, expand_as, tile, unstack

* fix unique unstack dim 0

* Update expand_v2_op.cc

* fix unique_op format
---
 paddle/fluid/operators/unique_op.h  | 5 ++++-
 paddle/fluid/operators/unstack_op.h | 2 +-
 python/paddle/fluid/layers/nn.py    | 2 ++
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/paddle/fluid/operators/unique_op.h b/paddle/fluid/operators/unique_op.h
index 99793ecd244..66b0543771f 100644
--- a/paddle/fluid/operators/unique_op.h
+++ b/paddle/fluid/operators/unique_op.h
@@ -403,7 +403,10 @@ class UniqueKernel : public framework::OpKernel<T> {
     bool return_index = context.Attr<bool>("return_index");
     bool return_inverse = context.Attr<bool>("return_inverse");
     bool return_counts = context.Attr<bool>("return_counts");
-
+    if (x->numel() == 0) {
+      out->mutable_data<T>(context.GetPlace());
+      return;
+    }
     if (axis_vec.empty()) {
       framework::VisitDataTypeTiny(
           data_type,
diff --git a/paddle/fluid/operators/unstack_op.h b/paddle/fluid/operators/unstack_op.h
index 82118b69270..cfd4d6bce83 100644
--- a/paddle/fluid/operators/unstack_op.h
+++ b/paddle/fluid/operators/unstack_op.h
@@ -149,7 +149,7 @@ class UnStackKernel : public framework::OpKernel<T> {
       dx_datas[i] = dx[i]->mutable_data<T>(ctx.GetPlace());
     }
     auto dy_data = dy->data<T>();
-
+    if (dy->numel() == 0) return;
     int pre = 1;
     for (int i = 0; i < axis; ++i) pre *= dy->dims()[i];
     int total_num = dy->numel();
diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py
index 515d4a5c0ef..75b0392ab6a 100755
--- a/python/paddle/fluid/layers/nn.py
+++ b/python/paddle/fluid/layers/nn.py
@@ -10315,6 +10315,8 @@ def unstack(x, axis=0, num=None):
     if in_dygraph_mode():
         if num == None:
             num = x.shape[axis]
+        if num == 0:
+            return []
         return _C_ops.unstack(x, num, 'axis', int(axis), 'num', num)
 
     helper = LayerHelper('unstack', **locals())
-- 
GitLab