Skip to content

Commit c157be0

Browse files
kshyattkshyatt
committed
Try letting Enzyme make rules for insert/remove unit
1 parent 67e57a0 commit c157be0

5 files changed

Lines changed: 96 additions & 116 deletions

File tree

‎ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ module TensorKitEnzymeExt
33
using Enzyme
44
using TensorKit
55
import TensorKit as TK
6-
using TensorKit: subblock
6+
using TensorKit: subblock, block
77
using VectorInterface
88
using TensorOperations: TensorOperations, IndexTuple, Index2Tuple, linearize
99
import TensorOperations as TO

‎ext/TensorKitEnzymeExt/indexmanipulations.jl‎

Lines changed: 0 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -148,92 +148,3 @@ function EnzymeRules.reverse(
148148
end
149149
return (nothing, nothing)
150150
end
151-
152-
for insertunit in (:insertleftunit, :insertrightunit)
153-
@eval begin
154-
function EnzymeRules.augmented_primal(
155-
config::EnzymeRules.RevConfigWidth{1},
156-
func::Const{typeof($insertunit)},
157-
::Type{RT},
158-
tsrc::Annotation{<:AbstractTensorMap},
159-
ival::Const{<:Val};
160-
kwargs...
161-
) where {RT}
162-
if tsrc.val isa TensorMap && !get(kwargs, :copy, false) && !isa(tsrc, Const)
163-
tsrc_cache = tsrc.val
164-
tdst = $insertunit(tsrc.val, ival.val; kwargs...)
165-
Δtdst = $insertunit(tsrc.dval, ival.val; kwargs...)
166-
else
167-
tsrc_cache = nothing
168-
tdst = $insertunit(tsrc.val, ival.val; kwargs...)
169-
Δtdst = make_zero(tdst)
170-
end
171-
primal = EnzymeRules.needs_primal(config) ? tdst : nothing
172-
shadow = EnzymeRules.needs_shadow(config) ? Δtdst : nothing
173-
cache = (tsrc_cache, tdst, Δtdst)
174-
return EnzymeRules.AugmentedReturn(primal, shadow, cache)
175-
end
176-
function EnzymeRules.reverse(
177-
config::EnzymeRules.RevConfigWidth{1},
178-
func::Const{typeof($insertunit)},
179-
::Type{RT},
180-
cache,
181-
tsrc::Annotation{<:AbstractTensorMap},
182-
ival::Const{<:Val};
183-
kwargs...
184-
) where {RT}
185-
tsrc_cache, tdst, Δtdst = cache
186-
# note: since data is already shared for <:TensorMap, don't have to do anything here!
187-
if isnothing(tsrc_cache) && !isa(tsrc, Const)
188-
for (c, b) in blocks(Δtdst)
189-
add!(block(tsrc.dval, c), b)
190-
end
191-
end
192-
return (nothing, nothing)
193-
end
194-
end
195-
end
196-
197-
function EnzymeRules.augmented_primal(
198-
config::EnzymeRules.RevConfigWidth{1},
199-
func::Const{typeof(removeunit)},
200-
::Type{RT},
201-
tsrc::Annotation{<:AbstractTensorMap},
202-
ival::Const{<:Val};
203-
kwargs...
204-
) where {RT}
205-
# tdst shares data with tsrc if <:TensorMap & copy=false, in this case we have to deal with correctly
206-
# sharing address spaces
207-
if tsrc.val isa TensorMap && !get(kwargs, :copy, false) && !isa(tsrc, Const)
208-
tsrc_cache = tsrc.val
209-
tdst = removeunit(tsrc.val, ival.val; kwargs...)
210-
Δtdst = removeunit(tsrc.dval, ival.val)
211-
else
212-
tsrc_cache = nothing
213-
tdst = removeunit(tsrc.val, ival.val; kwargs...)
214-
Δtdst = make_zero(tdst)
215-
end
216-
primal = EnzymeRules.needs_primal(config) ? tdst : nothing
217-
shadow = EnzymeRules.needs_shadow(config) ? Δtdst : nothing
218-
cache = (tsrc_cache, tdst, Δtdst)
219-
return EnzymeRules.AugmentedReturn(primal, shadow, cache)
220-
end
221-
222-
function EnzymeRules.reverse(
223-
config::EnzymeRules.RevConfigWidth{1},
224-
func::Const{typeof(removeunit)},
225-
::Type{RT},
226-
cache,
227-
tsrc::Annotation{<:AbstractTensorMap},
228-
ival::Const{<:Val};
229-
kwargs...
230-
) where {RT}
231-
tsrc_cache, tdst, Δtdst = cache
232-
# note: since data for <: TensorMap is already shared, don't have to do anything here!
233-
if isnothing(tsrc_cache) && !isa(tsrc, Const)
234-
for (c, b) in blocks(Δtdst)
235-
add!(block(tsrc.dval, c), b)
236-
end
237-
end
238-
return (nothing, nothing)
239-
end

‎ext/TensorKitEnzymeExt/utility.jl‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,39 @@ function EnzymeRules.reverse(
5252
return (nothing, nothing)
5353
end
5454

55+
function EnzymeRules.augmented_primal(
56+
config::EnzymeRules.RevConfigWidth{1},
57+
::Const{typeof(block)},
58+
::Type{RT},
59+
t::Annotation{<:AbstractTensorMap},
60+
c::Annotation{<:Sector},
61+
) where {RT}
62+
ret = EnzymeRules.needs_primal(config) ? block(t.val, c.val) : nothing
63+
dret = if !isa(t, Const) && EnzymeRules.needs_shadow(config)
64+
block(t.dval, c.val)
65+
elseif EnzymeRules.needs_shadow(config)
66+
Enzyme.make_zero(ret)
67+
else
68+
nothing
69+
end
70+
return EnzymeRules.AugmentedReturn(ret, dret, dret)
71+
end
72+
73+
function EnzymeRules.reverse(
74+
config::EnzymeRules.RevConfigWidth{1},
75+
::Const{typeof(block)},
76+
::Type{RT},
77+
cache,
78+
t::Annotation{<:AbstractTensorMap},
79+
c::Annotation{<:Sector},
80+
) where {RT}
81+
dret = cache
82+
if !isnothing(dret) && !isa(t, Const)
83+
block(t.dval, c.val) .= dret
84+
end
85+
return (nothing, nothing)
86+
end
87+
5588
function EnzymeRules.forward(
5689
config::EnzymeRules.FwdConfigWidth{1},
5790
::Const{typeof(subblock)},
@@ -78,9 +111,38 @@ function EnzymeRules.forward(
78111
end
79112
end
80113

114+
function EnzymeRules.forward(
115+
config::EnzymeRules.FwdConfigWidth{1},
116+
::Const{typeof(block)},
117+
::Type{RT},
118+
t::Annotation{<:AbstractTensorMap},
119+
c::Annotation{<:Sector},
120+
) where {RT}
121+
ret = EnzymeRules.needs_primal(config) ? block(t.val, c.val) : nothing
122+
dret = if !isa(t, Const) && EnzymeRules.needs_shadow(config)
123+
block(t.dval, c.val)
124+
elseif EnzymeRules.needs_shadow(config)
125+
Enzyme.make_zero(ret)
126+
else
127+
nothing
128+
end
129+
if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config)
130+
return Duplicated(ret, dret)
131+
elseif EnzymeRules.needs_primal(config)
132+
return ret
133+
elseif EnzymeRules.needs_shadow(config)
134+
return dret
135+
else
136+
return nothing
137+
end
138+
end
139+
81140
@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any) = nothing
82141
@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any, ::Any) = nothing
83142
@inline EnzymeRules.inactive(::typeof(TensorKit.artin_braid), ::Any, ::Any) = nothing
143+
@inline EnzymeRules.inactive(::typeof(TensorKit.insertleftunit), ::HomSpace, ::Any) = nothing
144+
@inline EnzymeRules.inactive(::typeof(TensorKit.insertrightunit), ::HomSpace, ::Any) = nothing
145+
@inline EnzymeRules.inactive(::typeof(TensorKit.removeunit), ::HomSpace, ::Any) = nothing
84146
@inline EnzymeRules.inactive(::typeof(TensorKit.sectorstructure), ::Any) = nothing
85147
@inline EnzymeRules.inactive(::typeof(TensorKit.degeneracystructure), ::Any) = nothing
86148
@inline EnzymeRules.inactive(::typeof(TensorKit.select), s::HomSpace, i::Index2Tuple) = nothing

‎test/enzyme-indexmanipulations-unit/insertunit.jl‎

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,27 @@ using Random
66
spacelist = ad_spacelist(fast_tests)
77
eltypes = (Float64, ComplexF64)
88

9-
if VERSION > v"1.11.0-rc" # https://github.com/QuantumKitHub/TensorKit.jl/issues/457
10-
@timedtestset verbose = true "Enzyme - Index Manipulations (insertunit):" begin
11-
@timedtestset verbose = true "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, TA in (Duplicated,)
12-
atol = default_tol(T)
13-
rtol = default_tol(T)
14-
A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
15-
@testset for insertunit in (insertleftunit, insertrightunit)
16-
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol)
17-
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(4), Const); atol, rtol)
18-
EnzymeTestUtils.test_reverse(insertunit, TA, (A', TA), (Val(2), Const); atol, rtol)
19-
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol, fkwargs = (copy = false,))
20-
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(2), Const); atol, rtol, fkwargs = (copy = true,))
21-
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(3), Const); atol, rtol, fkwargs = (copy = false, dual = true, conj = true))
22-
EnzymeTestUtils.test_reverse(insertunit, TA, (A', TA), (Val(3), Const); atol, rtol, fkwargs = (copy = false, dual = true, conj = true))
23-
end
9+
@timedtestset verbose = true "Enzyme - Index Manipulations (insertunit):" begin
10+
@timedtestset verbose = true "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, TA in (Duplicated,)
11+
atol = default_tol(T)
12+
rtol = default_tol(T)
13+
A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
14+
@testset for insertunit in (insertleftunit, insertrightunit)
15+
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol)
16+
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(4), Const); atol, rtol)
17+
EnzymeTestUtils.test_reverse(insertunit, TA, (A', TA), (Val(2), Const); atol, rtol)
18+
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol, fkwargs = (copy = false,))
19+
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(2), Const); atol, rtol, fkwargs = (copy = true,))
20+
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(3), Const); atol, rtol, fkwargs = (copy = false, dual = true, conj = true))
21+
EnzymeTestUtils.test_reverse(insertunit, TA, (A', TA), (Val(3), Const); atol, rtol, fkwargs = (copy = false, dual = true, conj = true))
22+
23+
EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol)
24+
EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(4), Const); atol, rtol)
25+
EnzymeTestUtils.test_forward(insertunit, TA, (A', TA), (Val(2), Const); atol, rtol)
26+
EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol, fkwargs = (copy = false,))
27+
EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(2), Const); atol, rtol, fkwargs = (copy = true,))
28+
EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(3), Const); atol, rtol, fkwargs = (copy = false, dual = true, conj = true))
29+
EnzymeTestUtils.test_forward(insertunit, TA, (A', TA), (Val(3), Const); atol, rtol, fkwargs = (copy = false, dual = true, conj = true))
2430
end
2531
end
2632
end

‎test/enzyme-indexmanipulations-unit/removeunit.jl‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,18 @@ using Random
66
spacelist = ad_spacelist(fast_tests)
77
eltypes = (Float64, ComplexF64)
88

9-
if VERSION > v"1.11.0-rc" # https://github.com/QuantumKitHub/TensorKit.jl/issues/457
10-
@timedtestset verbose = true "Enzyme - Index Manipulations (removeunit):" begin
11-
@timedtestset verbose = true "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, TB in (Duplicated,)
12-
atol = default_tol(T)
13-
rtol = default_tol(T)
14-
A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
15-
for i in 1:2
16-
B = insertleftunit(A, i; dual = rand(Bool))
17-
EnzymeTestUtils.test_reverse(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = false,))
18-
EnzymeTestUtils.test_reverse(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = true,))
19-
end
9+
@timedtestset verbose = true "Enzyme - Index Manipulations (removeunit):" begin
10+
@timedtestset verbose = true "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, TB in (Duplicated,)
11+
atol = default_tol(T)
12+
rtol = default_tol(T)
13+
A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
14+
for i in 1:2
15+
B = insertleftunit(A, i; dual = rand(Bool))
16+
EnzymeTestUtils.test_reverse(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = false,))
17+
EnzymeTestUtils.test_reverse(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = true,))
18+
19+
EnzymeTestUtils.test_forward(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = false,))
20+
EnzymeTestUtils.test_forward(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = true,))
2021
end
2122
end
2223
end

0 commit comments

Comments
 (0)