diff --git a/src/Transformations/ReplaceArraysInIO/replace_dist_arrays_in_io.cpp b/src/Transformations/ReplaceArraysInIO/replace_dist_arrays_in_io.cpp
index 02c7035..c4da117 100644
--- a/src/Transformations/ReplaceArraysInIO/replace_dist_arrays_in_io.cpp
+++ b/src/Transformations/ReplaceArraysInIO/replace_dist_arrays_in_io.cpp
@@ -17,6 +17,11 @@ using std::pair;
 
 #define DEBUG_TRACE 0
 
+static inline bool isArrayDeclaration(SgStatement* st)
+{
+    return isSgDeclarationStatement(st) || st->variant() == DIM_STAT || st->variant() == COMM_STAT;
+}
+
 static SgExpression* findExprWithVariant(SgExpression* exp, int variant)
 {
     if (exp)
@@ -644,18 +649,17 @@ void replaceDistributedArraysInIO(vector<ParallelRegion*>& regions,
                                                                newDeclsToInclude, 
                                                                copied_syms);
                             
-
                                 // original declaration statement
                                 auto *decl_stmt = SgStatement::getStatementByFileAndLine(decl_place.first, decl_place.second);
 
-                                if (!decl_stmt || !isSgDeclarationStatement(decl_stmt))
+                                if (!decl_stmt || !isArrayDeclaration(decl_stmt))
                                     printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
 
                                 auto dynamic_array = checkAssumedShape(decl_stmt, array_name);  
 
                                 while (decl_stmt && !isSgExecutableStatement(decl_stmt))
                                 {
-                                    if (isSgDeclarationStatement(decl_stmt) &&
+                                    if (isArrayDeclaration(decl_stmt) &&
                                         decl_stmt->expr(0) &&
                                         decl_stmt->expr(0)->lhs() &&
                                         decl_stmt->expr(0)->lhs()->symbol() == copied_symbol.second)
@@ -667,7 +671,7 @@ void replaceDistributedArraysInIO(vector<ParallelRegion*>& regions,
                                 }
                                 
                                 // created declaration statement
-                                if (!decl_stmt || !isSgDeclarationStatement(decl_stmt))
+                                if (!decl_stmt || !isArrayDeclaration(decl_stmt))
                                     printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
 
                                 // insert !$SPF ANALYSIS(PROCESS_PRIVATE(array)) directive before declaration statement