diff --git a/example/taylor_expansion.jl b/example/taylor_expansion.jl
index c198f5ac..97e8fc3f 100644
--- a/example/taylor_expansion.jl
+++ b/example/taylor_expansion.jl
@@ -1,4 +1,5 @@
 using FeynmanDiagram
+using FeynmanDiagram.Taylor
 using FeynmanDiagram.ComputationalGraphs:
     eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative, count_operation
 using FeynmanDiagram.Utility:
@@ -24,25 +25,29 @@ using FeynmanDiagram.Taylor:
 # set_variables("x y", order=3)
 # @time T5 = taylorexpansion!(G5)
 # print(T5)
-set_variables("x", order=3)
+set_variables("x y z a", order=7)
 @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
+# 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)
-    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")
+    @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/FeynmanDiagram.jl b/src/FeynmanDiagram.jl
index e429cb29..8d25dbc8 100644
--- a/src/FeynmanDiagram.jl
+++ b/src/FeynmanDiagram.jl
@@ -128,7 +128,7 @@ export optimize!, optimize, merge_all_chains!, merge_all_linear_combinations!, r
 include("TaylorSeries/TaylorSeries.jl")
 using .Taylor
 export Taylor
-export TaylorSeries, set_variables, taylor_factorial, getcoeff, getderivative
+export TaylorSeries, set_variables, taylor_factorial, getcoeff
 
 
 include("backend/compiler.jl")
@@ -181,7 +181,7 @@ export evalNaive, showTree
 include("utility.jl")
 using .Utility
 export Utility
-export taylorexpansion!, build_derivative_backAD!
+export taylorexpansion!
 ##################### precompile #######################
 # precompile as the final step of the module definition:
 if ccall(:jl_generating_output, Cint, ()) == 1   # if we're precompiling the package
diff --git a/src/TaylorSeries/TaylorSeries.jl b/src/TaylorSeries/TaylorSeries.jl
index 5fb84015..2dd12b15 100644
--- a/src/TaylorSeries/TaylorSeries.jl
+++ b/src/TaylorSeries/TaylorSeries.jl
@@ -26,7 +26,7 @@ export get_order, get_numvars,
     get_variable_names, get_variable_symbols,
     # jacobian, hessian, jacobian!, hessian!,
     displayBigO, use_show_default,
-    getcoeff, getderivative, taylor_factorial
+    getcoeff, taylor_factorial
 # function __init__()
 #     @static if !isdefined(Base, :get_extension)
 #         @require IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253" begin
diff --git a/src/TaylorSeries/arithmetic.jl b/src/TaylorSeries/arithmetic.jl
index 08924aad..5e73b852 100644
--- a/src/TaylorSeries/arithmetic.jl
+++ b/src/TaylorSeries/arithmetic.jl
@@ -33,7 +33,7 @@ function Base.:*(c1::Number, g2::TaylorSeries{T}) where {T}
 end
 
 """
-    function Base.:+(g1::TaylorSeries{T,V}, g2::TaylorSeries{T,V}) where {T,V}
+    function Base.:+(g1::TaylorSeries{T}, g2::TaylorSeries{T}) where {T}
 
     Returns a taylor series `g1 + g2` representing the addition of `g2` with `g1`.
 
@@ -57,13 +57,13 @@ end
 
 
 """
-    function Base.:+(g1::TaylorSeries{T,V}, g2::TaylorSeries{T,V}) where {T,V}
-
-    Returns a taylor series `g1 + g2` representing the addition of `g2` with `g1`.
+    function Base.:+(g1::TaylorSeries{T}, c::S) where {T,S<:Number}
 
+    Returns a taylor series `g1 + c` representing the addition of constant `c` with `g1`.
+    
 # Arguments:
-- `g1`  First taylor series
-- `g2`  Second taylor series
+- `g1`  Taylor series
+- `c`  Constant
 """
 function Base.:+(g1::TaylorSeries{T}, c::S) where {T,S<:Number}
     g = TaylorSeries{T}()
@@ -79,6 +79,16 @@ function Base.:+(g1::TaylorSeries{T}, c::S) where {T,S<:Number}
     return g
 end
 
+
+"""
+    function Base.:+(c::S, g1::TaylorSeries{T}) where {S<:Number,T}
+
+    Returns a taylor series `g1 + c` representing the addition of constant `c` with `g1`.
+    
+# Arguments:
+- `g1`  Taylor series
+- `c`  Constant
+"""
 function Base.:+(c::S, g1::TaylorSeries{T}) where {S<:Number,T}
     g = TaylorSeries{T}()
     g.coeffs = copy(g1.coeffs)
@@ -110,7 +120,15 @@ function Base.:-(g1::TaylorSeries{T}, c::S) where {T,S<:Number}
     return g1 + (-1 * c)
 end
 
+"""
+    function taylor_binomial(o1::Array{Int,1}, o2::Array{Int,1})
+
+    Return the taylor binomial prefactor when product two high-order derivatives with order o1 and o2.
 
+    # Arguments:
+    - `o1`  Order of first derivative
+    - `o2`  Order of second derivative
+"""
 function taylor_binomial(o1::Array{Int,1}, o2::Array{Int,1})
     @assert length(o1) == length(o2)
     result = 1
@@ -123,6 +141,15 @@ function taylor_binomial(o1::Array{Int,1}, o2::Array{Int,1})
     return result
 end
 
+
+"""
+    function taylor_factorial(o::Array{Int,1})
+
+    Return the taylor factorial prefactor with order o.
+
+    # Arguments:
+    - `o`  Order of the taylor coefficient
+"""
 function taylor_factorial(o::Array{Int,1})
     result = 1
     for i in eachindex(o)
@@ -130,8 +157,9 @@ function taylor_factorial(o::Array{Int,1})
     end
     return result
 end
+
 """
-    function Base.:*(g1::TaylorSeries{T,V}, g2::TaylorSeries{T,V}) where {T,V}
+    function Base.:*(g1::TaylorSeries{T}, g2::TaylorSeries{T}) where {T}
 
     Returns a taylor series `g1 * g2` representing the product of `g2` with `g1`.
 
@@ -159,34 +187,66 @@ function Base.:*(g1::TaylorSeries{T}, g2::TaylorSeries{T}) where {T}
     return g
 end
 
-# function findidx(a::TaylorSeries, o::Array{Int,1})
-#     @assert length(o) == get_numvars()
-#     return findfirst(isequal(o), a.order)
-# end
 
+"""
+    function getcoeff(g::TaylorSeries, order::Array{Int,1})
+
+    Return the taylor coefficients with given order in taylor series g.
+
+# Arguments:
+- `g`  Taylor series
+- `order`  Order of target coefficients
+"""
 function getcoeff(g::TaylorSeries, order::Array{Int,1})
     if haskey(g.coeffs, order)
         return g.coeffs[order]
-        #return 1 / taylor_factorial(order) * g.coeffs[order]
     else
         return nothing
     end
 end
 
+"""
+    function getderivative(g::TaylorSeries, order::Array{Int,1})
+
+    Return the derivative with given order in taylor series g.
+
+# Arguments:
+- `g`  Taylor series
+- `order`  Order of derivative
+"""
 function getderivative(g::TaylorSeries, order::Array{Int,1})
     if haskey(g.coeffs, order)
-        #return g.coeffs[order]
         return taylor_factorial(order) * g.coeffs[order]
     else
         return nothing
     end
 end
 
-function Base.one(x::TaylorSeries{T}) where {T}
-    g = TaylorSeries{T}()
-    g.coeffs[zeros(Int, get_numvars)] = one(T)
-    return g
+
+"""
+    function  Base.one(g::TaylorSeries{T}) where {T}
+
+    Return a constant one for a given taylor series.
+
+# Arguments:
+- `g`  Taylor series
+"""
+function Base.one(g::TaylorSeries{T}) where {T}
+    unity = TaylorSeries{T}()
+    unity.coeffs[zeros(Int, get_numvars)] = one(T)
+    return unity
 end
+
+
+"""
+    function Base.:^(x::TaylorSeries, p::Integer)
+
+    Return the power of taylor series x^p, where p is an integer. 
+
+# Arguments:
+- `x`  Taylor series
+- 'p' Power index
+"""
 function Base.:^(x::TaylorSeries, p::Integer)
     p == 1 && return copy(x)
     p == 0 && return one(x)
@@ -198,6 +258,8 @@ end
 function square(x::TaylorSeries)
     return x * x
 end
+
+# power_by_squaring; slightly modified from base/intfuncs.jl
 function power_by_squaring(x::TaylorSeries, p::Integer)
     p == 1 && return copy(x)
     p == 0 && return one(x)
diff --git a/src/TaylorSeries/parameter.jl b/src/TaylorSeries/parameter.jl
index 62f9cd96..cf243872 100644
--- a/src/TaylorSeries/parameter.jl
+++ b/src/TaylorSeries/parameter.jl
@@ -1,7 +1,8 @@
 """
     ParamsTaylor
 
-DataType holding the current parameters for `TaylorSeries`. This part of code 
+    DataType holding the current parameters for `TaylorSeries`. 
+    This part of code is adopted from TaylorSeries.jl (https://github.com/JuliaDiff/TaylorSeries.jl)
 
 **Fields:**
 
@@ -39,11 +40,10 @@ end
 """
     set_variables([T::Type], names::String; [order=get_order(), numvars=-1])
 
-Return a `Taylor{T}` vector with each entry representing an
+Return a `TaylorSeries{T}` vector with each entry representing an
 independent variable. `names` defines the output for each variable
 (separated by a space). The default type `T` is `Float64`,
 and the default for `order` is the one defined globally.
-Changing the `order` or `numvars` resets the hash_tables.
 
 If `numvars` is not specified, it is inferred from `names`. If only
 one variable name is defined and `numvars>1`, it uses this name with
diff --git a/src/TaylorSeries/print.jl b/src/TaylorSeries/print.jl
index 83227b3a..17abac49 100644
--- a/src/TaylorSeries/print.jl
+++ b/src/TaylorSeries/print.jl
@@ -1,4 +1,4 @@
-# This file is part of the TaylorSeries.jl Julia package, MIT license
+#This part of code is adopted from https://github.com/JuliaDiff/TaylorSeries.jl
 
 # Control the display of the big 𝒪 notation
 const bigOnotation = Bool[true]
diff --git a/src/utility.jl b/src/utility.jl
index 031232c5..6eb0e5e2 100644
--- a/src/utility.jl
+++ b/src/utility.jl
@@ -1,18 +1,26 @@
 module Utility
 using ..ComputationalGraphs
-using ..ComputationalGraphs: Sum, Prod, Power
+using ..ComputationalGraphs: Sum, Prod, Power, decrement_power
 using ..ComputationalGraphs: build_all_leaf_derivative, eval!
 using ..ComputationalGraphs.AbstractTrees
+
 using ..Taylor
 
 
 """
-    function graph_to_taylor(graph::G, varidx::Array{Int,1}=collect(1:get_numvars())) where {G<:Graph}
+    function graph_to_taylor(graph::G, var::Array{Bool,1}=fill(true, get_numvars())) where {G<:Graph}
+
+    Return a taylor series of graph g. If not provided, by default, assume that g depends on all variables. 
+
+#Arguments
+
+- `graph` Target graph. Must be a leaf.
+- `var` The variables graph depends on.
 """
-function graph_to_taylor(graph::G, varidx::Array{Int,1}=collect(1:get_numvars())) where {G<:Graph}
+function graph_to_taylor(graph::G, var::Array{Bool,1}=fill(true, get_numvars())) where {G<:Graph}
     @assert isleaf(graph)
     maxorder = get_order()
-    ordtuple = ((idx in varidx) ? (0:maxorder) : (0:0) for idx in 1:get_numvars())
+    ordtuple = ((var[idx]) ? (0:maxorder) : (0:0) for idx in 1:get_numvars())
     result = TaylorSeries{G}()
     for order in collect(Iterators.product(ordtuple...)) #varidx specifies the variables graph depends on. Iterate over all taylor coefficients of those variables.
         o = collect(order)
@@ -24,9 +32,20 @@ function graph_to_taylor(graph::G, varidx::Array{Int,1}=collect(1:get_numvars())
     return result
 end
 
-
-function graph_to_taylor_withmap!(g::G; coeffmode=true, varidx::Array{Int,1}=collect(1:get_numvars()), chainrule_map_leaf::Dict{Int,Dict{Int,G}}=Dict{Int,Dict{Int,G}}()) where {G<:Graph}
+"""
+    graph_to_taylor_withmap!(g::G; coeffmode=true, var::Array{Int,1}=collect(1:get_numvars())) where {G<:Graph}
+    
+    Return a taylor series of graph g, together with a map of chain relation ship between generated derivatives.
+    This function is only internally used for constructing high order derivatives by naive nested forward AD.
+    It is only for banch mark purpose and not exported.
+# Arguments:
+- `g`  Target graph 
+- `coeffmode` If true, the generated taylor series saves taylor coefficietnts with the factorial prefactor. If false, the taylor series saves derivatives instead
+- `var` The index of variables graph depends on
+"""
+function graph_to_taylor_withmap!(g::G; coeffmode=true, var::Array{Bool,1}=fill(true, get_numvars())) where {G<:Graph}
     @assert isleaf(g)
+    chainrule_map_leaf = Dict{Int,Dict{Int,G}}()
     maxorder = get_order()
     current_func = Dict(zeros(Int, get_numvars()) => g)
     result = TaylorSeries{G}()
@@ -38,40 +57,51 @@ function graph_to_taylor_withmap!(g::G; coeffmode=true, varidx::Array{Int,1}=col
             if !haskey(chainrule_map_leaf, func.id)
                 chainrule_map_leaf[func.id] = Dict{Int,G}()
             end
-            for idx in varidx
-                ordernew = copy(order)
-                ordernew[idx] += 1
-                if !haskey(result.coeffs, ordernew)
-                    if coeffmode
-                        funcAD = Graph([]; operator=Sum(), factor=g.factor)
+            for idx in eachindex(var)
+                if var[idx]
+                    ordernew = copy(order)
+                    ordernew[idx] += 1
+                    if !haskey(result.coeffs, ordernew)
+                        if coeffmode
+                            funcAD = Graph([]; operator=Sum(), factor=g.factor)
+                        else
+                            #funcAD = taylor_factorial(ordernew) * Graph([]; operator=Sum(), factor=g.factor)
+                            funcAD = Graph([]; operator=Sum(), factor=taylor_factorial(ordernew) * g.factor)
+                        end
+                        new_func[ordernew] = funcAD
+                        result.coeffs[ordernew] = funcAD
+                        chainrule_map_leaf[func.id][idx] = funcAD
                     else
-                        #funcAD = taylor_factorial(ordernew) * Graph([]; operator=Sum(), factor=g.factor)
-                        funcAD = Graph([]; operator=Sum(), factor=taylor_factorial(ordernew) * g.factor)
+                        chainrule_map_leaf[func.id][idx] = result.coeffs[ordernew]
                     end
-                    new_func[ordernew] = funcAD
-                    result.coeffs[ordernew] = funcAD
-                    chainrule_map_leaf[func.id][idx] = funcAD
-                else
-                    chainrule_map_leaf[func.id][idx] = result.coeffs[ordernew]
                 end
             end
         end
         current_func = new_func
     end
 
-    return result
+    return result, chainrule_map_leaf
 end
 
 @inline apply(::Type{Sum}, diags::Vector{T}, factors::Vector{F}) where {T<:TaylorSeries,F<:Number} = sum(d * f for (d, f) in zip(diags, factors))
 @inline apply(::Type{Prod}, diags::Vector{T}, factors::Vector{F}) where {T<:TaylorSeries,F<:Number} = prod(d * f for (d, f) in zip(diags, factors))
 @inline apply(::Type{Power{N}}, diags::Vector{T}, factors::Vector{F}) where {N,T<:TaylorSeries,F<:Number} = (diags[1])^N * factors[1]
 
+"""
+    function taylorexpansion!(graph::G, taylormap::Dict{Int,T}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph,T<:TaylorSeries}
+    
+    Return the taylor Series of a graph. If taylor series of the leaves of this graph is not provided, by default we assume the leaves depend on all variables.
+    
+# Arguments:
+- `graph`  Target graph 
+- `taylormap` The taylor series corresponding to each node of graphs. The taylor series of leafs can be provided as input
+- `varidx` The index of variables graph depends on
+"""
 function taylorexpansion!(graph::G, taylormap::Dict{Int,T}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph,T<:TaylorSeries}
     if isempty(taylormap)
         for g in Leaves(graph)
             if !haskey(taylormap, g.id)
-                #taylormap[g.id] = graph_to_taylor(g)
-                taylormap[g.id] = graph_to_taylor_withmap!(g)
+                taylormap[g.id] = graph_to_taylor(g)
             end
         end
     end
@@ -86,19 +116,19 @@ function taylorexpansion!(graph::G, taylormap::Dict{Int,T}=Dict{Int,TaylorSeries
     return taylormap[rootid]
 end
 
+
+#Functions below generate high order derivatives with naive nested forward AD. This part would be significantly refactored later with 
+# Taylor Series API.
+
 function build_derivative_backAD!(g::G, leaftaylor::Dict{Int,TaylorSeries{G}}=Dict{Int,TaylorSeries{G}}()) where {G<:Graph}
     chainrule_map_leaf = Dict{Int,Dict{Int,G}}()
     for leaf in Leaves(g)
         if !haskey(leaftaylor, leaf.id)
-            leaftaylor[leaf.id] = graph_to_taylor_withmap!(leaf; coeffmode=false, chainrule_map_leaf=chainrule_map_leaf)
+            leaftaylor[leaf.id], map = graph_to_taylor_withmap!(leaf; coeffmode=false)
+            chainrule_map_leaf = merge(chainrule_map_leaf, map)
         end
     end
-    for (i, leaf) in chainrule_map_leaf
-        for (j, coeff) in leaf
-            print("test $(i) $(j) $(coeff) $(eval!(coeff))\n")
-        end
-        print("\n")
-    end
+
     leafAD, chainrule_map = build_all_leaf_derivative(g)
     current_func = Dict(zeros(Int, get_numvars()) => g)
 
diff --git a/test/taylor.jl b/test/taylor.jl
index f7acb68f..f785294d 100644
--- a/test/taylor.jl
+++ b/test/taylor.jl
@@ -20,5 +20,26 @@ using FeynmanDiagram: Taylor as Taylor
     @test getcoeff(F1, [1, 2, 0, 0, 0]) == 3.0
     @test getcoeff(F1, [3, 0, 0, 0, 0]) == 1.0
     @test getcoeff(F1, [0, 3, 0, 0, 0]) == 1.0
+    using FeynmanDiagram.ComputationalGraphs:
+        eval!, forwardAD, node_derivative, backAD, build_all_leaf_derivative, count_operation
+    using FeynmanDiagram.Utility:
+        taylorexpansion!, build_derivative_backAD!
+    g1 = Graph([])
+    g2 = Graph([])
+    g3 = Graph([], factor=2.0)
+    G3 = g1
+    G4 = 1.0 * g1 * g1
+    G5 = 1.0 * (3.0 * G3 + 0.5 * G4)
+    G6 = (1.0 * g1 + 2.0 * g2) * (g1 + g3)
+    using FeynmanDiagram.Taylor:
+        TaylorSeries, getcoeff, set_variables
 
+    set_variables("x y z", order=5)
+    for G in [G3, G4, G5, G6]
+        T = 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
\ No newline at end of file