Skip to content

Commit 832ceb7

Browse files
kshyattkshyatt
authored
Enzyme doesn't need a custom rule for trace_permute (#480)
* Does Enzyme need a custom rule for trace_permute * Working
1 parent bb8bf9d commit 832ceb7

3 files changed

Lines changed: 63 additions & 90 deletions

File tree

‎ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ module TensorKitEnzymeExt
33
using Enzyme
44
using TensorKit
55
import TensorKit as TK
6+
using TensorKit: subblock
67
using VectorInterface
78
using TensorOperations: TensorOperations, IndexTuple, Index2Tuple, linearize
89
import TensorOperations as TO

‎ext/TensorKitEnzymeExt/tensoroperations.jl‎

Lines changed: 0 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -116,93 +116,3 @@ function EnzymeRules.forward(
116116
return nothing
117117
end
118118
end
119-
120-
# tensortrace!
121-
# ------------
122-
123-
function EnzymeRules.augmented_primal(
124-
config::EnzymeRules.RevConfigWidth{1},
125-
func::Const{typeof(TensorKit.trace_permute!)},
126-
::Type{RT},
127-
C::Annotation{<:AbstractTensorMap},
128-
A::Annotation{<:AbstractTensorMap},
129-
p::Const{<:Index2Tuple},
130-
q::Const{<:Index2Tuple},
131-
α::Annotation{<:Number},
132-
β::Annotation{<:Number},
133-
backend::Const,
134-
) where {RT}
135-
C_cache = !isa(β, Const) ? copy(C.val) : nothing
136-
A_cache = EnzymeRules.overwritten(config)[3] ? copy(A.val) : nothing
137-
At = if !isa(α, Const)
138-
At = TO.tensortrace(A.val, p.val, q.val, false, One(), backend.val)
139-
add!(C.val, At, α.val, β.val)
140-
At
141-
else
142-
TensorKit.trace_permute!(C.val, A.val, p.val, q.val, α.val, β.val, backend.val)
143-
nothing
144-
end
145-
primal = EnzymeRules.needs_primal(config) ? C.val : nothing
146-
shadow = EnzymeRules.needs_shadow(config) ? C.dval : nothing
147-
cache = (C_cache, A_cache, At)
148-
return EnzymeRules.AugmentedReturn(primal, shadow, cache)
149-
end
150-
151-
152-
function EnzymeRules.reverse(
153-
config::EnzymeRules.RevConfigWidth{1},
154-
func::Const{typeof(TensorKit.trace_permute!)},
155-
::Type{RT},
156-
cache,
157-
C::Annotation{<:AbstractTensorMap},
158-
A::Annotation{<:AbstractTensorMap},
159-
p::Const{<:Index2Tuple},
160-
q::Const{<:Index2Tuple},
161-
α::Annotation{<:Number},
162-
β::Annotation{<:Number},
163-
backend::Const,
164-
) where {RT}
165-
C_cache, A_cache, At = cache
166-
Aval = something(A_cache, A.val)
167-
Cval = something(C_cache, C.val)
168-
!isa(A, Const) && !isa(C, Const) && TensorKit.trace_permute_pullback_ΔA!(A.dval, C.dval, Aval, p.val, q.val, α.val, backend.val)
169-
Δαr = pullback_dα(α, C, At)
170-
Δβr = pullback_dβ(β, C, Cval)
171-
!isa(C, Const) && pullback_dC!(C.dval, β.val)
172-
return nothing, nothing, nothing, nothing, Δαr, Δβr, nothing
173-
end
174-
175-
function EnzymeRules.forward(
176-
config::EnzymeRules.FwdConfigWidth{1},
177-
func::Const{typeof(TensorKit.trace_permute!)},
178-
::Type{RT},
179-
C::Annotation{<:AbstractTensorMap},
180-
A::Annotation{<:AbstractTensorMap},
181-
p::Annotation{<:Index2Tuple},
182-
q::Annotation{<:Index2Tuple},
183-
α::Annotation{<:Number},
184-
β::Annotation{<:Number},
185-
backend::Const,
186-
) where {RT}
187-
# dD = dα * tr(A) + α * tr(dA) + dβ * C + β * dC
188-
# dC1 = dβ * C + β * dC
189-
if !isa(C, Const)
190-
if isa(β, Const)
191-
scale!(C.dval, β.val)
192-
else
193-
add!(C.dval, C.val, β.dval, β.val)
194-
end
195-
!isa(α, Const) && TensorKit.trace_permute!(C.dval, A.val, p.val, q.val, α.dval, One(), backend.val)
196-
!isa(A, Const) && TensorKit.trace_permute!(C.dval, A.dval, p.val, q.val, α.val, One(), backend.val)
197-
end
198-
TensorKit.trace_permute!(C.val, A.val, p.val, q.val, α.val, β.val, backend.val)
199-
if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config)
200-
return C
201-
elseif EnzymeRules.needs_primal(config)
202-
return C.val
203-
elseif EnzymeRules.needs_shadow(config)
204-
return C.dval
205-
else
206-
return nothing
207-
end
208-
end

‎ext/TensorKitEnzymeExt/utility.jl‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,68 @@ pullback_dC!(ΔC, β::Number) = scale!(ΔC, conj(β))
1919
@inline EnzymeRules.inactive_type(::Type{<:TensorKit.GenericTreeTransformer}) = true
2020
@inline EnzymeRules.inactive_type(::Type{<:TensorKit.VectorSpace}) = true
2121

22+
function EnzymeRules.augmented_primal(
23+
config::EnzymeRules.RevConfigWidth{1},
24+
::Const{typeof(subblock)},
25+
::Type{RT},
26+
t::Annotation{<:AbstractTensorMap},
27+
f::Annotation,
28+
) where {RT}
29+
ret = EnzymeRules.needs_primal(config) ? subblock(t.val, f.val) : nothing
30+
dret = if !isa(t, Const) && EnzymeRules.needs_shadow(config)
31+
subblock(t.dval, f.val)
32+
elseif EnzymeRules.needs_shadow(config)
33+
Enzyme.make_zero(ret)
34+
else
35+
nothing
36+
end
37+
return EnzymeRules.AugmentedReturn(ret, dret, dret)
38+
end
39+
40+
function EnzymeRules.reverse(
41+
config::EnzymeRules.RevConfigWidth{1},
42+
::Const{typeof(subblock)},
43+
::Type{RT},
44+
cache,
45+
t::Annotation{<:AbstractTensorMap},
46+
f::Annotation,
47+
) where {RT}
48+
dret = cache
49+
if !isnothing(dret) && !isa(t, Const)
50+
subblock(t.dval, f.val) .= dret
51+
end
52+
return (nothing, nothing)
53+
end
54+
55+
function EnzymeRules.forward(
56+
config::EnzymeRules.FwdConfigWidth{1},
57+
::Const{typeof(subblock)},
58+
::Type{RT},
59+
t::Annotation{<:AbstractTensorMap},
60+
f::Annotation,
61+
) where {RT}
62+
ret = EnzymeRules.needs_primal(config) ? subblock(t.val, f.val) : nothing
63+
dret = if !isa(t, Const) && EnzymeRules.needs_shadow(config)
64+
subblock(t.dval, f.val)
65+
elseif EnzymeRules.needs_shadow(config)
66+
Enzyme.make_zero(ret)
67+
else
68+
nothing
69+
end
70+
if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config)
71+
return Duplicated(ret, dret)
72+
elseif EnzymeRules.needs_primal(config)
73+
return ret
74+
elseif EnzymeRules.needs_shadow(config)
75+
return dret
76+
else
77+
return nothing
78+
end
79+
end
80+
81+
@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any) = nothing
82+
@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any, ::Any) = nothing
83+
@inline EnzymeRules.inactive(::typeof(TensorKit.artin_braid), ::Any, ::Any) = nothing
2284
@inline EnzymeRules.inactive(::typeof(TensorKit.sectorstructure), ::Any) = nothing
2385
@inline EnzymeRules.inactive(::typeof(TensorKit.degeneracystructure), ::Any) = nothing
2486
@inline EnzymeRules.inactive(::typeof(TensorKit.select), s::HomSpace, i::Index2Tuple) = nothing

0 commit comments

Comments
 (0)