diff --git a/Project.toml b/Project.toml
index f5f2ac4..6a478e1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "YaoToEinsum"
 uuid = "9b173c7b-dc24-4dc5-a0e1-adab2f7b6ba9"
 authors = ["GiggleLiu <cacate0129@gmail.com> and contributors"]
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -10,14 +10,13 @@ Yao = "5872b779-8223-5990-8dd0-5abbb0748c8c"
 
 [compat]
 OMEinsum = "0.6"
-Yao = "0.6"
+Yao = "0.7"
 julia = "1"
 
 [extras]
 OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-YaoExtensions = "7a06699c-c960-11e9-3c98-9f78548b5f0f"
 SymEngine = "123dc426-2d89-5057-bbad-38513e3affd8"
 
 [targets]
-test = ["Test", "OMEinsumContractionOrders", "YaoExtensions", "SymEngine"]
+test = ["Test", "OMEinsumContractionOrders", "SymEngine"]
diff --git a/src/YaoToEinsum.jl b/src/YaoToEinsum.jl
index 0921a5a..5fc1a0a 100644
--- a/src/YaoToEinsum.jl
+++ b/src/YaoToEinsum.jl
@@ -24,7 +24,7 @@ function EinBuilder(::Type{T}, n::Int) where T
 end
 newlabel!(eb::EinBuilder) = (eb.maxlabel[] += 1; eb.maxlabel[])
 
-function add_gate!(eb::EinBuilder{T}, b::PutBlock{N,C}) where {T, N,C}
+function add_gate!(eb::EinBuilder{T}, b::PutBlock{D,C}) where {T,D,C}
     return add_matrix!(eb, C, mat(T, b.content), collect(b.locs))
 end
 # general and diagonal gates
@@ -39,7 +39,7 @@ function add_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector)
     return eb
 end
 # swap gate
-function add_gate!(eb::EinBuilder{T}, b::PutBlock{N,2,ConstGate.SWAPGate}) where {T,N}
+function add_gate!(eb::EinBuilder{T}, b::PutBlock{2,2,ConstGate.SWAPGate}) where {T,N}
     lj = eb.slots[b.locs[2]]
     eb.slots[b.locs[2]] = eb.slots[b.locs[1]]
     eb.slots[b.locs[1]] = lj
@@ -47,7 +47,7 @@ function add_gate!(eb::EinBuilder{T}, b::PutBlock{N,2,ConstGate.SWAPGate}) where
 end
 
 # control gates
-function add_gate!(eb::EinBuilder{T}, b::ControlBlock{N,BT,C,M}) where {T, N,BT,C,M}
+function add_gate!(eb::EinBuilder{T}, b::ControlBlock{BT,C,M}) where {T, BT,C,M}
     return add_controlled_matrix!(eb, M, mat(T, b.content), collect(b.locs), collect(b.ctrl_locs), collect(b.ctrl_config))
 end
 function add_controlled_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector, control_locs, control_vals) where T
@@ -138,19 +138,19 @@ julia> yao2einsum(c; initial_state=Dict(1=>0, 2=>1), final_state=Dict(1=>ArrayRe
 (1, 2, 4∘2, 5∘1, 6∘3∘5, 5, 4 -> 6∘3, AbstractArray{ComplexF64}[[1.0 + 0.0im, 0.0 + 0.0im], [0.0 + 0.0im, 1.0 + 0.0im], [0.0 + 0.0im 1.0 + 0.0im; 1.0 + 0.0im 0.0 + 0.0im], [0.0 + 0.0im 0.0 - 1.0im; 0.0 + 1.0im 0.0 + 0.0im], [1.0 + 0.0im 0.0 + 0.0im; 0.0 + 0.0im 1.0 + 0.0im;;; 0.0 + 0.0im 0.0 - 1.0im; 0.0 + 1.0im 0.0 + 0.0im], [0.6 + 0.0im, 0.0 + 0.8im], [0.0 + 0.0im, 1.0 + 0.0im]])
 ```
 """
-function yao2einsum(circuit::AbstractBlock; initial_state::Dict=Dict{Int,Int}(), final_state::Dict=Dict{Int,Int}())
+function yao2einsum(circuit::AbstractBlock{D}; initial_state::Dict=Dict{Int,Int}(), final_state::Dict=Dict{Int,Int}()) where {D}
     T = promote_type(ComplexF64, dict_regtype(initial_state), dict_regtype(final_state), Yao.parameters_eltype(circuit))
-    vec_initial_state = Dict{Int,ArrayReg{1,T}}([k=>render_single_qubit_state(T, v) for (k, v) in initial_state])
-    vec_final_state = Dict{Int,ArrayReg{1,T}}([k=>render_single_qubit_state(T, v) for (k, v) in final_state])
+    vec_initial_state = Dict{Int,ArrayReg{D,T}}([k=>render_single_qubit_state(T, v) for (k, v) in initial_state])
+    vec_final_state = Dict{Int,ArrayReg{D,T}}([k=>render_single_qubit_state(T, v) for (k, v) in final_state])
     yao2einsum(circuit, vec_initial_state, vec_final_state)
 end
 dict_regtype(d::Dict) = promote_type(_regtype.(values(d))...)
-_regtype(::ArrayReg{1,VT}) where {VT} = VT
+_regtype(::ArrayReg{D,VT}) where {D,VT} = VT
 _regtype(::Int) = ComplexF64
 render_single_qubit_state(::Type{T}, x::Int) where T = x == 0 ? zero_state(T, 1) : product_state(T, bit"1")
-render_single_qubit_state(::Type{T}, x::ArrayReg{1}) where T = ArrayReg(collect(T, statevec(x)))
+render_single_qubit_state(::Type{T}, x::ArrayReg) where T = ArrayReg(collect(T, statevec(x)))
 
-function yao2einsum(circuit::AbstractBlock, initial_state::Dict{Int,<:ArrayReg{1,T}}, final_state::Dict{Int,<:ArrayReg{1,T}}) where T
+function yao2einsum(circuit::AbstractBlock{D}, initial_state::Dict{Int,<:ArrayReg{D,T}}, final_state::Dict{Int,<:ArrayReg{D,T}}) where {D,T}
     n = nqubits(circuit)
     eb = EinBuilder(T, n)
     openindices = Int[]
diff --git a/test/runtests.jl b/test/runtests.jl
index 3ec0b3b..2cf9e81 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,7 +1,7 @@
 using YaoToEinsum
 using Test, OMEinsum, OMEinsumContractionOrders
 using Yao
-using YaoExtensions: qft_circuit, variational_circuit, rand_google53
+using Yao.EasyBuild: qft_circuit, variational_circuit, rand_google53
 using SymEngine
 
 @testset "YaoToEinsum.jl" begin