Fix Cached_op with static_shape=true (#15298)

* Fix

* run ci

Author: ZhennanQin

src/imperative/cached_op.cc

@@ -81,6 +81,7 @@ struct CachedOp::CachedOpState {
 
   std::vector<NDArray> buff;
   std::vector<NDArray*> arrays;
+  std::vector<NDArray*> arrays_with_in_out;
   std::vector<OpReqType> array_reqs;
 
   std::vector<OpStatePtr> op_states;
@@ -762,7 +763,8 @@ OpStatePtr CachedOp::StaticForward(
   // We are going to add input and output arrays to the array list.
   // The input and output arrays should only be valid for this run,
   // so we shouldn't modify the state's array list.
-  auto arrays = state.arrays;
+  state.arrays_with_in_out = state.arrays;
+  auto& arrays = state.arrays_with_in_out;
   if (config_.static_shape) {
     for (auto i : config_.param_indices) {
       auto nid = idx.input_nodes()[i];
@@ -1063,7 +1065,8 @@ void CachedOp::StaticBackward(
   // We are going to add input and output arrays to the array list.
   // The input and output arrays should only be valid for this run,
   // so we shouldn't modify the state's array list.
-  auto arrays = state.arrays;
+  state.arrays_with_in_out = state.arrays;
+  auto& arrays = state.arrays_with_in_out;
   for (size_t i = 0; i < state.info.bwd_input_eid.size(); ++i) {
     auto eid = state.info.bwd_input_eid[i];
     if (eid == kEidNotExist) {

src/nnvm/legacy_op_util.cc

@@ -79,7 +79,6 @@ class OperatorState {
  public:
   OperatorState(Operator *opr, const OperatorProperty *prop) {
     opr_ = opr;
-    fwd_init_ = bwd_init_ = false;
 
     in_data_fwd_.resize(prop->ListArguments().size());
     in_data_bwd_.resize(prop->ListArguments().size());
@@ -110,47 +109,39 @@ class OperatorState {
                const std::vector<TBlob>& inputs,
                const std::vector<OpReqType>& req,
                const std::vector<TBlob>& outputs) {
-    if (!fwd_init_) {
-      CHECK_EQ(inputs.size(), in_data_fwd_.size() + aux_data_.size());
-      CHECK_EQ(outputs.size(), out_data_.size());
-      // in_data_bwd_ has the same tblobs as the ones in in_data_fwd_, except that the ones
-      // referred by arg_data_ptr_ will be overriden
-      for (size_t i = 0; i < in_data_fwd_.size(); ++i) in_data_fwd_[i] = inputs[i];
-      for (size_t i = 0; i < in_data_fwd_.size(); ++i) in_data_bwd_[i] = inputs[i];
-      for (size_t i = 0; i < aux_data_.size(); ++i) {
-        aux_data_[i] = inputs[i + in_data_fwd_.size()];
-      }
-      for (size_t i = 0; i < out_data_.size(); ++i) out_data_[i] = outputs[i];
-      fwd_init_ = true;
+    CHECK_EQ(inputs.size(), in_data_fwd_.size() + aux_data_.size());
+    CHECK_EQ(outputs.size(), out_data_.size());
+    // in_data_bwd_ has the same tblobs as the ones in in_data_fwd_, except that the ones
+    // referred by arg_data_ptr_ will be overriden
+    for (size_t i = 0; i < in_data_fwd_.size(); ++i) in_data_fwd_[i] = inputs[i];
+    for (size_t i = 0; i < in_data_fwd_.size(); ++i) in_data_bwd_[i] = inputs[i];
+    for (size_t i = 0; i < aux_data_.size(); ++i) {
+      aux_data_[i] = inputs[i + in_data_fwd_.size()];
     }
+    for (size_t i = 0; i < out_data_.size(); ++i) out_data_[i] = outputs[i];
     opr_->Forward(ctx, in_data_fwd_, req, out_data_, aux_data_);
   }
 
   void Backward(const OpContext &ctx,
                 const std::vector<TBlob>& inputs,
                 const std::vector<OpReqType>& req,
                 const std::vector<TBlob>& outputs) {
-    if (!bwd_init_) {
-      CHECK(fwd_init_);
-      CHECK_EQ(arg_data_ptr_.size() + aux_data_.size(), inputs.size());
-      // override tblobs pointed by arg_data_ptr_ since they might not contain
-      // initialized data during forward pass.
-      for (size_t i = 0; i < arg_data_ptr_.size(); ++i) {
-        *arg_data_ptr_[i] = inputs[i];
-      }
-      for (size_t i = 0; i < aux_data_.size(); ++i) {
-        aux_data_[i] = inputs[inputs.size() - aux_data_.size() + i];
-      }
-      CHECK_EQ(outputs.size(), in_grad_.size());
-      for (size_t i = 0; i < outputs.size(); ++i) in_grad_[i] = outputs[i];
-      bwd_init_ = true;
+    CHECK_EQ(arg_data_ptr_.size() + aux_data_.size(), inputs.size());
+    // override tblobs pointed by arg_data_ptr_ since they might not contain
+    // initialized data during forward pass.
+    for (size_t i = 0; i < arg_data_ptr_.size(); ++i) {
+      *arg_data_ptr_[i] = inputs[i];
+    }
+    for (size_t i = 0; i < aux_data_.size(); ++i) {
+      aux_data_[i] = inputs[inputs.size() - aux_data_.size() + i];
     }
+    CHECK_EQ(outputs.size(), in_grad_.size());
+    for (size_t i = 0; i < outputs.size(); ++i) in_grad_[i] = outputs[i];
     opr_->Backward(ctx, out_grad_, in_data_bwd_, out_data_, req, in_grad_, aux_data_);
   }
 
  private:
   Operator *opr_;
-  bool fwd_init_, bwd_init_;
   // input data blobs for forward and backward
   // in_data_fwd_ and in_data_bwd_ will hold different tblobs when StorageFallbackOpExecutor
   // performs storage fallback on a non-default input NDArray. The one in in_data_fwd_ is