diff --git a/example/taylor_expansion.jl b/example/taylor_expansion.jl
index 2cfdb0af..4a0aa067 100644
--- a/example/taylor_expansion.jl
+++ b/example/taylor_expansion.jl
@@ -1,9 +1,32 @@
 using FeynmanDiagram
 using FeynmanDiagram.Taylor
 using FeynmanDiagram.ComputationalGraphs:
-    eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative, count_operation
+    eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative
 using FeynmanDiagram.Utility:
-    taylorexpansion, build_derivative_backAD!
+    taylorexpansion!, build_derivative_backAD!, count_operation
+
+function benchmark_AD(glist::Vector{T}) where {T<:Graph}
+    taylormap = Dict{Int,TaylorSeries{T}}()
+    totaloperation = [0, 0]
+    taylorlist = Vector{TaylorSeries{T}}()
+    for g in glist
+        @time t, taylormap = taylorexpansion!(g; taylormap=taylormap)
+
+
+        operation = count_operation(t)
+        totaloperation = totaloperation + operation
+        push!(taylorlist, t)
+        print("operation number: $(operation)\n")
+        t_compare = build_derivative_backAD!(g)
+        for (order, coeff) in (t_compare.coeffs)
+            @assert (eval!(coeff)) == (eval!(Taylor.taylor_factorial(order) * t.coeffs[order]))
+        end
+    end
+
+    total_uniqueoperation = count_operation(taylorlist)
+    print(" total operation number: $(length(taylorlist)) $(totaloperation) $(total_uniqueoperation)\n")
+    return total_uniqueoperation
+end
 g1 = Graph([])
 g2 = Graph([])
 g3 = Graph([]) #, factor=2.0)
@@ -13,40 +36,15 @@ g6 = Graph([])
 G3 = g1
 G4 = 1.0 * g1 * g2
 G5 = 1.0 * (3.0 * G3 + 0.5 * G4)
-#G6 = (0.5 * g1 * g2 + 0.5 * g2 * g3)
-#G6 = (g1 + g2) * (0.5 * g1 + g3) * g1 #   (0.5 * g1 + g3)
-#G6 = g1 * g2 * g3 * g4 * g5 * g6
 G6 = (1.0 * g1 + 2.0 * g2) * (g1 + g3)
-#G6 = 1.5 * g1*g1 + 0.5 * g2 * 1.5 * g1 +    0.5*g2*g3
+
 using FeynmanDiagram.Taylor:
     TaylorSeries, getcoeff, set_variables
 
+set_variables("x y", orders=[3, 2])
+
+benchmark_AD([G3, G4, G5, G6])
 
-# set_variables("x y", order=3)
-# @time T5 = taylorexpansion(G5)
-# print(T5)
-set_variables("x y", order=[3, 2])
-#set_variables("x y z a", order=[1, 2, 3, 2])
-@time T5 = taylorexpansion(G6)
-#order = [0, 0, 0, 0, 0, 0]
-#@time print(T5.coeffs[order])
-print("$(count_operation(T5.coeffs))\n")
-# for (order, coeff) in (T5.coeffs)
-#     #gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
-#     #println("$(order)  ", gs, "\n")
-#     print("$(order) $(eval!(coeff)) $(eval!(getcoeff(T5,order))) $(coeff.id) $(count_operation(coeff))\n")
-# end
-
-print("TaylorSeries $(T5)\n")
-
-@time T5_compare = build_derivative_backAD!(G6)
-print("$(count_operation(T5_compare.coeffs))\n")
-for (order, coeff) in (T5_compare.coeffs)
-    @assert (eval!(coeff)) == (eval!(Taylor.taylor_factorial(order) * T5.coeffs[order]))
-    # gs = Compilers.to_julia_str([coeff,], name="eval_graph!")
-    # println("$(order)  ", gs, "\n")
-    # print("$(order) $(eval!(coeff)) $(eval!(getderivative(T5,order))) $(count_operation(coeff))\n")
-end
 
 
 
diff --git a/src/computational_graph/tree_properties.jl b/src/computational_graph/tree_properties.jl
index b315e13d..82f4cce4 100644
--- a/src/computational_graph/tree_properties.jl
+++ b/src/computational_graph/tree_properties.jl
@@ -156,7 +156,7 @@ function count_operation(g::Array{G}) where {G<:Graph}
 end
 
 
-function count_operation(g::Dict{Array{Int,1},G}) where {G<:Graph}
+function count_operation(g::Dict{Vector{Int},G}) where {G<:Graph}
     visited = Set{Int}()
     totalsum = 0
     totalprod = 0
diff --git a/src/frontend/diagtree.jl b/src/frontend/diagtree.jl
index e0c5c43f..fa246b65 100644
--- a/src/frontend/diagtree.jl
+++ b/src/frontend/diagtree.jl
@@ -55,4 +55,26 @@ function Graph!(d::DiagTree.Diagram{W}; map=Dict{Int,DiagTree.DiagramId}()) wher
     map[tree.id] = d.id
 
     return root, map
+end
+
+"""
+    function extract_var_dependence(map::Dict{Int,DiagTree.DiagramId}, ::Type{ID}, numvars::Int)
+    
+    Given a map between graph id and DiagramId, extract the variable dependence of all graphs.
+
+# Arguments:
+- `map::Dict{Int,DiagTree.DiagramId}`:  A dictionary mapping graph ids to DiagramIds. DiagramId stores the diagram information of the corresponding graph. 
+- `ID`:  The particular type of ID that has the given variable dependence. 
+- `numvars`: The number of variables which the diagram depends on.
+"""
+function extract_var_dependence(map::Dict{Int,DiagTree.DiagramId}, ::Type{ID}; numvars::Int=1) where {ID<:PropagatorId}
+    var_dependence = Dict{Int,Vector{Bool}}()
+    for (id, diagID) in map
+        if diagID isa ID
+            var_dependence[id] = [true for _ in 1:numvars]
+        else
+            var_dependence[id] = [false for _ in 1:numvars]
+        end
+    end
+    return var_dependence
 end
\ No newline at end of file
diff --git a/src/utility.jl b/src/utility.jl
index 3deea045..70ac437f 100644
--- a/src/utility.jl
+++ b/src/utility.jl
@@ -2,24 +2,18 @@ module Utility
 using ..ComputationalGraphs
 using ..ComputationalGraphs: Sum, Prod, Power, decrement_power
 using ..ComputationalGraphs: build_all_leaf_derivative, eval!
+import ..ComputationalGraphs: count_operation
 using ..ComputationalGraphs.AbstractTrees
 
 using ..Taylor
 
 
-"""
-    function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
-
-    Return a taylor series of graph g. If variable dependence is not specified, by default, assume all leaves of graph depend on all variables. 
+# Internal function that performs taylor expansion on a single graph node recursively.
+function taylorexpansion!(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
+    if haskey(taylormap, graph.id) #If already exist, use taylor series in taylormap.
+        return taylormap[graph.id], taylormap
 
-#Arguments
-
-- `graph` Target graph.
-- `var_dependence::Dict{Int,Vector{Bool}}` The variables graph leaves depend on. Should map each leaf ID of g to a Vector{Bool},
-    indicating the taylor variables it depends on. By default, assumes all leaves depend on all variables.
-"""
-function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
-    if isleaf(graph)
+    elseif isleaf(graph)
         maxorder = get_orders()
         if haskey(var_dependence, graph.id)
             var = var_dependence[graph.id]
@@ -33,27 +27,41 @@ function taylorexpansion(graph::G, var_dependence::Dict{Int,Vector{Bool}}=Dict{I
             coeff = Graph([]; operator=Sum(), factor=graph.factor)
             result.coeffs[o] = coeff
         end
-        return result
+        taylormap[graph.id] = result
+        return result, taylormap
     else
-        taylormap = Dict{Int,TaylorSeries{G}}() #Saves the taylor series corresponding to each nodes of the graph
-        for g in Leaves(graph)
-            if !haskey(taylormap, g.id)
-                taylormap[g.id] = taylorexpansion(g, var_dependence)
-            end
-        end
+        taylormap[graph.id] = apply(graph.operator, [taylorexpansion!(sub, var_dependence; taylormap=taylormap)[1] for sub in graph.subgraphs], graph.subgraph_factors)
+        return taylormap[graph.id], taylormap
+    end
+end
 
-        rootid = -1
-        for g in PostOrderDFS(graph) # postorder traversal will visit all subdiagrams of a diagram first
-            rootid = g.id
-            if isleaf(g) || haskey(taylormap, g.id)
-                continue
-            end
-            taylormap[g.id] = apply(g.operator, [taylormap[sub.id] for sub in g.subgraphs], g.subgraph_factors)
-        end
-        return taylormap[rootid]
+
+"""
+    function taylorexpansion!(graph::G, taylormap::Dict{Int,TaylorSeries{G}}(), var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}()) where {G<:Graph}
+
+    Return a taylor series of graph g. If variable dependence is not specified, by default, assume all leaves of graph depend on all variables. The taylor series of all nodes of graph is 
+    saved in the taylormap dictionary.
+
+#Arguments
+
+- `graph` Target graph.
+- `var_dependence::Dict{Int,Vector{Bool}}` The variables graph leaves depend on. Should map each leaf ID of g to a Vector{Bool},
+    indicating the taylor variables it depends on. By default, assumes all leaves depend on all variables.
+- `taylormap::Dict{Int,TaylorSeries{G}}` The taylor series correponding to graph nodes. Should map each graph node ID (does not need to belong to input graph) to a taylor series.
+    All new taylor series generated by taylor expansion will be added into the expansion.
+"""
+
+
+function taylorexpansion!(graphs::Vector{G}, var_dependence::Dict{Int,Vector{Bool}}=Dict{Int,Vector{Bool}}(); taylormap::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
+    result = Vector{TaylorSeries{G}}()
+    for graph in graphs
+        taylor, _ = taylorexpansion!(graph, var_dependence; taylormap=taylormap)
+        push!(result, taylor)
     end
+    return result, taylormap
 end
 
+
 """
     taylorexpansion_withmap(g::G; coeffmode=true, var::Vector{Int}=collect(1:get_numvars())) where {G<:Graph}
     
@@ -225,4 +233,22 @@ function chainrule!(varidx::Int, dg::Array{G,1}, leaftaylor::Dict{Int,TaylorSeri
     end
 end
 
+function count_operation(g::TaylorSeries{G}) where {G<:Graph}
+    return count_operation(g.coeffs)
+end
+
+function count_operation(graphs::Vector{TaylorSeries{G}}) where {G<:Graph}
+    if length(graphs) == 0
+        return [0, 0]
+    else
+        allcoeffs = Vector{G}()
+        for g in graphs
+            for (order, coeffs) in g.coeffs
+                push!(allcoeffs, coeffs)
+            end
+        end
+        return count_operation(allcoeffs)
+    end
+end
+
 end
\ No newline at end of file
diff --git a/test/diagram_tree.jl b/test/diagram_tree.jl
index 3495b84c..b94f7a70 100644
--- a/test/diagram_tree.jl
+++ b/test/diagram_tree.jl
@@ -91,7 +91,7 @@ end
 
     DiagTree.uidreset()
     # We only consider the direct part of the above diagram
-    spin = 2.0
+    spin = 1.0
     D = 3
     kF, β, mass2 = 1.919, 0.5, 1.0
     Nk, Nt = 4, 2
@@ -107,13 +107,14 @@ end
     # plot_tree(droot_dg)
 
     DiagTree.eval!(root; eval=(x -> 1.0))
-    @test root.weight ≈ -2 + spin
+    factor = 1 / (2π)^D
+    @test root.weight ≈ (-2 + spin) * factor
 
     DiagTree.eval!(droot_dg; eval=(x -> 1.0))
-    @test root.weight ≈ (-2 + spin) * 2
+    @test droot_dg.weight ≈ (-2 + spin) * 2 * factor
 
     DiagTree.eval!(droot_dv; eval=(x -> 1.0))
-    @test root.weight ≈ (-2 + spin) * 2
+    @test droot_dv.weight ≈ (-2 + spin) * 2 * factor
 
     # #more sophisticated test of the weight evaluation
     varK = rand(D, Nk)
diff --git a/test/taylor.jl b/test/taylor.jl
index ce268535..3726f53f 100644
--- a/test/taylor.jl
+++ b/test/taylor.jl
@@ -23,7 +23,7 @@ using FeynmanDiagram: Taylor as Taylor
     using FeynmanDiagram.ComputationalGraphs:
         eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative, count_operation
     using FeynmanDiagram.Utility:
-        taylorexpansion, build_derivative_backAD!
+        taylorexpansion!, build_derivative_backAD!
     g1 = Graph([])
     g2 = Graph([])
     g3 = Graph([], factor=2.0)
@@ -34,10 +34,108 @@ using FeynmanDiagram: Taylor as Taylor
 
     set_variables("x y z", orders=[2, 3, 2])
     for G in [G3, G4, G5, G6]
-        T = taylorexpansion(G)
+        T, taylormap = taylorexpansion!(G)
         T_compare = build_derivative_backAD!(G)
         for (order, coeff) in T_compare.coeffs
             @test eval!(coeff) == eval!(taylor_factorial(order) * T.coeffs[order])
         end
     end
+
+end
+
+
+function getdiagram(spin=2.0, D=3, Nk=4, Nt=2)
+    """
+        k1-k3                     k2+k3 
+        |                         | 
+    t1.L ↑     t1.L       t2.L     ↑ t2.L
+        |-------------->----------|
+        |       |  k3+k4   |      |
+        |   v   |          |  v   |
+        |       |    k4    |      |
+        |--------------<----------|
+    t1.L ↑    t1.L        t2.L     ↑ t2.L
+        |                         | 
+        k1                        k2
+    """
+
+    DiagTree.uidreset()
+    # We only consider the direct part of the above diagram
+
+    paraG = DiagParaF64(type=GreenDiag,
+        innerLoopNum=0, totalLoopNum=Nk, loopDim=D,
+        hasTau=true, totalTauNum=Nt)
+    paraV = paraG
+
+    # #construct the propagator table
+    gK = [[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 0.0, 1.0]]
+    gT = [(1, 2), (2, 1)]
+    g = [Diagram{Float64}(BareGreenId(paraG, k=gK[i], t=gT[i]), name=:G) for i in 1:2]
+
+    vdK = [[0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 1.0, 0.0]]
+    # vdT = [[1, 1], [2, 2]]
+    vd = [Diagram{Float64}(BareInteractionId(paraV, ChargeCharge, k=vdK[i], permu=Di), name=:Vd) for i in 1:2]
+
+    veK = [[1, 0, -1, -1], [0, 1, 0, -1]]
+    # veT = [[1, 1], [2, 2]]
+    ve = [Diagram{Float64}(BareInteractionId(paraV, ChargeCharge, k=veK[i], permu=Ex), name=:Ve) for i in 1:2]
+
+    Id = GenericId(paraV)
+    # contruct the tree
+    ggn = Diagram{Float64}(Id, Prod(), [g[1], g[2]])
+    vdd = Diagram{Float64}(Id, Prod(), [vd[1], vd[2]], factor=spin)
+    vde = Diagram{Float64}(Id, Prod(), [vd[1], ve[2]], factor=-1.0)
+    ved = Diagram{Float64}(Id, Prod(), [ve[1], vd[2]], factor=-1.0)
+    vsum = Diagram{Float64}(Id, Sum(), [vdd, vde, ved])
+    root = Diagram{Float64}(Id, Prod(), [vsum, ggn], factor=1 / (2π)^D, name=:root)
+
+    return root, gK, gT, vdK, veK
+end
+
+@testset "Taylor AD of DiagTree" begin
+
+    DiagTree.uidreset()
+    # We only consider the direct part of the above diagram
+    spin = 0.5
+    D = 3
+    kF, β, mass2 = 1.919, 0.5, 1.0
+    Nk, Nt = 4, 2
+
+    root, gK, gT, vdK, veK = getdiagram(spin, D, Nk, Nt)
+
+    #optimize the diagram
+    DiagTree.optimize!([root,])
+
+    # autodiff
+    droot_dg = DiagTree.derivative([root,], BareGreenId)[1]
+    droot_dv = DiagTree.derivative([root,], BareInteractionId)[1]
+    # plot_tree(droot_dg)
+    factor = 1 / (2π)^D
+    DiagTree.eval!(root; eval=(x -> 1.0))
+    @test root.weight ≈ (-2 + spin) * factor
+
+    DiagTree.eval!(droot_dg; eval=(x -> 1.0))
+    @test droot_dg.weight ≈ (-2 + spin) * 2 * factor
+
+    DiagTree.eval!(droot_dv; eval=(x -> 1.0))
+    @test droot_dv.weight ≈ (-2 + spin) * 2 * factor
+
+    set_variables("x"; orders=[2])
+    g, map = FrontEnds.Graph!(root)
+    var_dependence = FrontEnds.extract_var_dependence(map, BareGreenId)
+    t, taylormap = taylorexpansion!(g, var_dependence)
+    order = [0]
+    @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor
+
+    order = [1]
+    @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor * 2 * taylor_factorial(order)
+
+    var_dependence = FrontEnds.extract_var_dependence(map, BareInteractionId)
+
+    t, taylormap = taylorexpansion!(g, var_dependence)
+    order = [0]
+    @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor
+
+    order = [1]
+    @test eval!(taylormap[g.id].coeffs[order]) ≈ (-2 + spin) * factor * 2 * taylor_factorial(order)
 end
\ No newline at end of file