TensorKit.jl/ext/TensorKitEnzymeExt/utility.jl at 61eba8a323b8ca89b737e4d9cac4592c8b22eb5d · QuantumKitHub/TensorKit.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
# Projection
# ----------
pullback_dα(α::Const, C::Const, A) = nothing
pullback_dα(α::Const, C::Annotation, A) = nothing
pullback_dα(α::Annotation, C::Const, A) = zero(α.val)
pullback_dα(α::Annotation, C::Annotation, A) = TK._pullback_dα(α.val, C.dval, A)

pullback_dβ(β::Const, C::Const, Ccache) = nothing
pullback_dβ(β::Const, C::Annotation, Ccache) = nothing
pullback_dβ(β::Annotation, C::Const, Ccache) = zero(β.val)
pullback_dβ(β::Annotation, C::Annotation, Ccache) = TK._pullback_dβ(β.val, C.dval, Ccache)

pullback_dC!(ΔC, β::Number) = scale!(ΔC, conj(β))

# Ignore derivatives
# ------------------

@inline EnzymeRules.inactive_type(::Type{<:TensorKit.FusionTree}) = true
@inline EnzymeRules.inactive_type(::Type{<:TensorKit.GenericTreeTransformer}) = true
@inline EnzymeRules.inactive_type(::Type{<:TensorKit.VectorSpace}) = true

function EnzymeRules.augmented_primal(
        config::EnzymeRules.RevConfigWidth{1},
        ::Const{typeof(subblock)},
        ::Type{RT},
        t::Annotation{<:AbstractTensorMap},
        f::Annotation,
    ) where {RT}
    ret = EnzymeRules.needs_primal(config) ? subblock(t.val, f.val) : nothing
    dret = if !isa(t, Const) && EnzymeRules.needs_shadow(config)
        subblock(t.dval, f.val)
    elseif EnzymeRules.needs_shadow(config)
        Enzyme.make_zero(ret)
    else
        nothing
    end
    return EnzymeRules.AugmentedReturn(ret, dret, dret)
end

function EnzymeRules.reverse(
        config::EnzymeRules.RevConfigWidth{1},
        ::Const{typeof(subblock)},
        ::Type{RT},
        cache,
        t::Annotation{<:AbstractTensorMap},
        f::Annotation,
    ) where {RT}
    dret = cache
    if !isnothing(dret) && !isa(t, Const)
        subblock(t.dval, f.val) .= dret
    end
    return (nothing, nothing)
end

function EnzymeRules.forward(
        config::EnzymeRules.FwdConfigWidth{1},
        ::Const{typeof(subblock)},
        ::Type{RT},
        t::Annotation{<:AbstractTensorMap},
        f::Annotation,
    ) where {RT}
    ret = EnzymeRules.needs_primal(config) ? subblock(t.val, f.val) : nothing
    dret = if !isa(t, Const) && EnzymeRules.needs_shadow(config)
        subblock(t.dval, f.val)
    elseif EnzymeRules.needs_shadow(config)
        Enzyme.make_zero(ret)
    else
        nothing
    end
    if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config)
        return Duplicated(ret, dret)
    elseif EnzymeRules.needs_primal(config)
        return ret
    elseif EnzymeRules.needs_shadow(config)
        return dret
    else
        return nothing
    end
end

@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any, ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.artin_braid), ::Any, ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.sectorstructure), ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.degeneracystructure), ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.select), s::HomSpace, i::Index2Tuple) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.flip), s::HomSpace, i::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.permute), s::HomSpace, i::Index2Tuple) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.braid), s::HomSpace, i::Index2Tuple, ::IndexTuple) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.compose), s1::HomSpace, s2::HomSpace) = nothing
@inline EnzymeRules.inactive(::typeof(TensorOperations.tensorcontract), c::HomSpace, p::Index2Tuple, α::Bool, b::HomSpace, q::Index2Tuple, β::Bool, pq::Index2Tuple) = nothing