Repository: intel/ipex-llm
Pull request: Clone graph keep their original prevs nodes and next nodes sequence
## What changes were proposed in this pull request?

Clone graph keep their original prevs nodes and next nodes order.
Current cloneGraph may change the next nodes order, thus will influence the update of tree-lstm (it has many shared component)

In the meantime, fix a bug in graph backward that caused by the following ut

"Graph backward" should "be correct when contains nested output"

## How was this patch tested?

unit test
test over tree lstm training

## Related links or issues (optional)
fixed https://github.com/intel-analytics/BigDL/issues/1743



Problem statement:
TreeLSTM example will throw exception and loss will be NaN
```
Exception in thread "main" java.lang.IllegalArgumentException: requirement failed: copyStatus:container's length mismatchexcepted 1, but get 0
        at scala.Predef$.require(Predef.scala:224)                                                                                                                                             
        at com.intel.analytics.bigdl.nn.Container.copyStatus(Container.scala:143)
        at com.intel.analytics.bigdl.nn.Container.copyStatus(Container.scala:40)                                                                                                               
        at com.intel.analytics.bigdl.nn.Container.copyStatus(Container.scala:148)           
        at com.intel.analytics.bigdl.nn.Container.copyStatus(Container.scala:40)                                                                                                               
        at com.intel.analytics.bigdl.nn.Utils$.copyModule(Utils.scala:238)
        at com.intel.analytics.bigdl.optim.DistriOptimizer.optimize(DistriOptimizer.scala:859)                                                                                                 
        at com.intel.analytics.bigdl.example.treeLSTMSentiment.Train$.train(Train.scala:124)
        at com.intel.analytics.bigdl.example.treeLSTMSentiment.Train$.main(Train.scala:40)                                                                                                     
        at com.intel.analytics.bigdl.example.treeLSTMSentiment.Train.main(Train.scala)
        at sun.reflect.NativeMethodAccessorImpl.invoke0(Native Method)                                                                                                                         
        at sun.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:62)
        at sun.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)                                                                                               
        at java.lang.reflect.Method.invoke(Method.java:497)
        at org.apache.spark.deploy.SparkSubmit$.org$apache$spark$deploy$SparkSubmit$$runMain(SparkSubmit.scala:729)                                                                            
        at org.apache.spark.deploy.SparkSubmit$.doRunMain$1(SparkSubmit.scala:185)
        at org.apache.spark.deploy.SparkSubmit$.submit(SparkSubmit.scala:210)                                                                                                                  
        at org.apache.spark.deploy.SparkSubmit$.main(SparkSubmit.scala:124)
        at org.apache.spark.deploy.SparkSubmit.main(SparkSubmit.scala)                                                 
```

Hints from issue discussion:

Commit message: Clone graph keep their original prevs nodes and next nodes sequence

Patch to review:
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/Graph.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/Graph.scala
index 03aaea59d6b..578781b850b 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/Graph.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/Graph.scala
@@ -509,14 +509,17 @@ class Graph[T: ClassTag](val inputs : Seq[ModuleNode[T]],
         activity.toTensor[T].add(other.toTensor[T])
       } else {
         val actTable = activity.toTable
-        val actLen = actTable.length()
         val otherTable = other.toTable
-        val otherLen = otherTable.length()
-        require(actLen == otherLen, "table length is not equal")
+        val maxActIndex = actTable.keySet.map(_.asInstanceOf[Int]).max
+        val maxOtherIndex = otherTable.keySet.map(_.asInstanceOf[Int]).max
+        val maxIndex = Math.max(maxActIndex, maxOtherIndex)
         var i = 1
-        while(i <= actLen) {
-          require(actTable[Activity](i) != null, "Invalid table element")
-          accActivity(actTable[Activity](i), otherTable[Activity](i))
+        while(i <= maxIndex) {
+          if (actTable.contains(i) && otherTable.contains(i)) {
+            accActivity(actTable[Activity](i), otherTable[Activity](i))
+          } else if (otherTable.contains(i)) {
+            actTable.insert(i, otherTable(i))
+          }
           i += 1
         }
         actTable
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/utils/DirectedGraph.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/utils/DirectedGraph.scala
index 17b3894d35a..3cd8b7553eb 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/utils/DirectedGraph.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/utils/DirectedGraph.scala
@@ -148,12 +148,17 @@ class DirectedGraph[T](val source : Node[T], val reverse : Boolean = false) exte
     })
     bfs.foreach(node => {
       if (reverseEdge) {
-        node.prevNodesAndEdges.foreach(prevNodeAndEdge => {
-          oldToNew.get(node).add(oldToNew.get(prevNodeAndEdge._1), prevNodeAndEdge._2)
+        node.nextNodesAndEdges.foreach(nextNodeAndEdge => {
+          // some node may be ignored
+          if (oldToNew.containsKey(nextNodeAndEdge._1)) {
+            oldToNew.get(nextNodeAndEdge._1).add(oldToNew.get(node), nextNodeAndEdge._2)
+          }
         })
       } else {
-        node.prevNodesAndEdges.foreach(prevNodeAndEdge => {
-          oldToNew.get(prevNodeAndEdge._1).add(oldToNew.get(node), prevNodeAndEdge._2)
+        node.nextNodesAndEdges.foreach(nextNodeAndEdge => {
+          if (oldToNew.containsKey(nextNodeAndEdge._1)) {
+            oldToNew.get(node).add(oldToNew.get(nextNodeAndEdge._1), nextNodeAndEdge._2)
+          }
         })
       }
     })
diff --git a/spark/dl/src/test/scala/com/intel/analytics/bigdl/utils/DirectedGraphSpec.scala b/spark/dl/src/test/scala/com/intel/analytics/bigdl/utils/DirectedGraphSpec.scala
index 1754cab9c87..7e909819865 100644
--- a/spark/dl/src/test/scala/com/intel/analytics/bigdl/utils/DirectedGraphSpec.scala
+++ b/spark/dl/src/test/scala/com/intel/analytics/bigdl/utils/DirectedGraphSpec.scala
@@ -315,7 +315,13 @@ class DirectedGraphSpec extends FlatSpec with Matchers {
 
     val graph = nodeA.graph()
     val cloneGraph = graph.cloneGraph()
-    graph.topologySort.map(_.element) should be(cloneGraph.topologySort.map(_.element))
+    val sort1 = graph.topologySort
+    val sort2 = cloneGraph.topologySort
+    sort1.map(_.element) should be(sort2.map(_.element))
+    sort1.zip(sort2).foreach(x => {
+      x._1.prevNodes.map(_.element) should be (x._2.prevNodes.map(_.element))
+      x._1.nextNodes.map(_.element) should be (x._2.nextNodes.map(_.element))
+    })
   }
 
   "Clone graph" should "should reuse the edge" in {


Code review comments:
diff_hunk: @@ -509,14 +509,17 @@ class Graph[T: ClassTag](val inputs : Seq[ModuleNode[T]],
         activity.toTensor[T].add(other.toTensor[T])
       } else {
         val actTable = activity.toTable
-        val actLen = actTable.length()
         val otherTable = other.toTable
-        val otherLen = otherTable.length()
-        require(actLen == otherLen, "table length is not equal")
+        val maxActIndex = actTable.keySet.map(_.asInstanceOf[Int]).max
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/Graph.scala
text: @user1:
try to use length instead of keySet.map

@user1:
otherTable.keySet.foreach?