diff --git a/lib/YaoBlocks/Project.toml b/lib/YaoBlocks/Project.toml
index 74816f70..3400f4b3 100644
--- a/lib/YaoBlocks/Project.toml
+++ b/lib/YaoBlocks/Project.toml
@@ -27,7 +27,7 @@ CacheServers = "0.2"
 ChainRulesCore = "1.11"
 DocStringExtensions = "0.8, 0.9"
 InteractiveUtils = "1"
-KrylovKit = "0.5, 0.6, 0.7"
+KrylovKit = "0.5, 0.6, 0.7, 0.8"
 LegibleLambdas = "0.2, 0.3"
 LinearAlgebra = "1"
 LuxurySparse = "0.7"
diff --git a/lib/YaoBlocks/src/primitive/rotation_gate.jl b/lib/YaoBlocks/src/primitive/rotation_gate.jl
index 80652aa1..09e6e591 100644
--- a/lib/YaoBlocks/src/primitive/rotation_gate.jl
+++ b/lib/YaoBlocks/src/primitive/rotation_gate.jl
@@ -137,7 +137,7 @@ cache_key(R::RotationGate) = R.theta
 
 iparams_range(::RotationGate{D,T,GT}) where {D,T,GT} = ((zero(T), T(2 * pi)),)
 
-occupied_locs(g::RotationGate) = occupied_locs(g.block)
+occupied_locs(g::RotationGate) = (1:nqudits(g)...,)
 
 function unsafe_getindex(::Type{T}, rg::RotationGate{D}, i::Integer, j::Integer) where {D,T}
     return (i==j ? cos(T(rg.theta)/2) : zero(T)) - im * sin(T(rg.theta)/2) * unsafe_getindex(T, rg.block, i, j)
diff --git a/lib/YaoBlocks/src/treeutils/optimise.jl b/lib/YaoBlocks/src/treeutils/optimise.jl
index 5ecbba8f..1f7412ae 100644
--- a/lib/YaoBlocks/src/treeutils/optimise.jl
+++ b/lib/YaoBlocks/src/treeutils/optimise.jl
@@ -7,7 +7,7 @@ using YaoBlocks: NotImplementedError
 
 include("to_basictypes.jl")
 
-export replace_block
+export replace_block, flatten_basic, simplify
 """
     replace_block(actor, tree::AbstractBlock) -> AbstractBlock
     replace_block(pair::Pair{Type{ST}, TT}, tree::AbstractBlock) -> AbstractBlock
@@ -208,6 +208,9 @@ export simplify
 
 const __default_simplification_rules__ =
     Function[merge_pauli, eliminate_nested, merge_scale, combine_similar]
+const __flatten_basic__ = Function[eliminate_nested, to_basictypes]
+
+flatten_basic(ex::AbstractBlock) = simplify(ex; rules = __flatten_basic__)
 
 # Inspired by MasonPotter/Symbolics.jl
 """
diff --git a/lib/YaoBlocks/src/treeutils/to_basictypes.jl b/lib/YaoBlocks/src/treeutils/to_basictypes.jl
index a59bce45..ec4ba19d 100644
--- a/lib/YaoBlocks/src/treeutils/to_basictypes.jl
+++ b/lib/YaoBlocks/src/treeutils/to_basictypes.jl
@@ -12,6 +12,7 @@ convert gates to basic types
 function to_basictypes end
 
 to_basictypes(block::PrimitiveBlock) = block
+to_basictypes(block::UnitaryChannel) = block
 function to_basictypes(block::AbstractBlock)
     throw(NotImplementedError(:to_basictypes, typeof(block)))
 end
diff --git a/lib/YaoBlocks/test/primitive/rotation_gate.jl b/lib/YaoBlocks/test/primitive/rotation_gate.jl
index 3f687439..97a01ce7 100644
--- a/lib/YaoBlocks/test/primitive/rotation_gate.jl
+++ b/lib/YaoBlocks/test/primitive/rotation_gate.jl
@@ -69,7 +69,7 @@ end
 
 @testset "occupied locs" begin
     g = rot(put(5, 2 => X), 0.5)
-    @test occupied_locs(g) == (2,)
+    @test occupied_locs(g) == (1,2,3,4,5)
 end
 
 @testset "instruct_get_element" begin
diff --git a/lib/YaoPlots/src/vizcircuit.jl b/lib/YaoPlots/src/vizcircuit.jl
index a054bce8..85dc3044 100644
--- a/lib/YaoPlots/src/vizcircuit.jl
+++ b/lib/YaoPlots/src/vizcircuit.jl
@@ -278,6 +278,10 @@ function draw!(c::CircuitGrid, p::Daggered{<:PrimitiveBlock}, address, controls)
     bts = length(controls)>=1 ? get_cbrush_texts(c, content(p)) : get_brush_texts(c, content(p))
     _draw!(c, [controls..., (getindex.(Ref(address), occupied_locs(p)), bts[1], bts[2]*"'")])
 end
+function draw!(c::CircuitGrid, p::UnitaryChannel, address, controls)
+    bts = (c.gatestyles.g, "*")
+    _draw!(c, [controls..., (getindex.(Ref(address), occupied_locs(p)), bts[1], bts[2])])
+end
 
 function draw!(c::CircuitGrid, p::Scale, address, controls)
     fp = YaoBlocks.factor(p)
diff --git a/lib/YaoPlots/test/vizcircuit.jl b/lib/YaoPlots/test/vizcircuit.jl
index d1bf25ec..aad1e179 100644
--- a/lib/YaoPlots/test/vizcircuit.jl
+++ b/lib/YaoPlots/test/vizcircuit.jl
@@ -86,4 +86,9 @@ end
     )
     YaoPlots.CircuitStyles.barrier_for_chain[] = true
     @test vizcircuit(circuit) isa Drawing
-end
\ No newline at end of file
+end
+
+@testset "rot igate" begin
+    @test plot(rot(igate(1), 1.)) isa Drawing
+    @test plot(rot(put(3, 1=>X), 1.)) isa Drawing
+end