BatchNLPKernels.jl/ext/BNKChainRulesCore.jl at 6e8cd0341e5da7138ffba36781f9b3ce4d413f0f · LearningToOptimize/BatchNLPKernels.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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
module BNKChainRulesCore

using BatchNLPKernels
using ChainRulesCore

function ChainRulesCore.rrule(::typeof(BatchNLPKernels.objective!), bm::BatchModel, X, Θ)
    y = BatchNLPKernels.objective!(bm, X, Θ)

    function obj_batch_pullback(Ȳ)
        Ȳ = ChainRulesCore.unthunk(Ȳ)
        gradients = BatchNLPKernels.objective_gradient!(bm, X, Θ)

        X̄ = gradients .* Ȳ'

        return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), X̄, ChainRulesCore.NoTangent()
    end

    return y, obj_batch_pullback
end
function ChainRulesCore.rrule(::typeof(BatchNLPKernels.objective!), bm::BatchModel, X)
    y = BatchNLPKernels.objective!(bm, X)

    function obj_batch_pullback(Ȳ)
        Ȳ = ChainRulesCore.unthunk(Ȳ)
        gradients = BatchNLPKernels.objective_gradient!(bm, X)

        X̄ = gradients .* Ȳ'

        return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), X̄
    end

    return y, obj_batch_pullback
end


function ChainRulesCore.rrule(::typeof(BatchNLPKernels.constraints!), bm::BatchModel, X, Θ)
    y = BatchNLPKernels.constraints!(bm, X, Θ)

    function cons_nln_batch_pullback(Ȳ)
        Ȳ = ChainRulesCore.unthunk(Ȳ)
        X̄ = BatchNLPKernels.constraints_jtprod!(bm, X, Θ, Ȳ)
        return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), X̄, ChainRulesCore.NoTangent()
    end

    return y, cons_nln_batch_pullback
end
function ChainRulesCore.rrule(::typeof(BatchNLPKernels.constraints!), bm::BatchModel, X)
    y = BatchNLPKernels.constraints!(bm, X)

    function cons_nln_batch_pullback(Ȳ)
        Ȳ = ChainRulesCore.unthunk(Ȳ)
        X̄ = BatchNLPKernels.constraints_jtprod!(bm, X, Ȳ)
        return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), X̄
    end

    return y, cons_nln_batch_pullback
end


function ChainRulesCore.rrule(::typeof(BatchNLPKernels._constraint_violations!), bm::BatchModel, V)
    Vc = BatchNLPKernels._constraint_violations!(bm, V)

    function constraint_violations_pullback(Ȳ)
        Ȳ = ChainRulesCore.unthunk(Ȳ)

        # violation(v, s) = max(s.l - v, v - s.u, 0)
        # ∂violation/∂v = -1 if v < s.l, +1 if v > s.u, 0 otherwise

        V̄ = if isempty(bm.viols_cons)
            zeros(eltype(V), size(V))
        else
            lower_viols = V .< bm.viols_cons.l
            upper_viols = V .> bm.viols_cons.u
            lower_viols .* (-Ȳ) .+ upper_viols .* Ȳ
        end

        return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), V̄
    end

    return Vc, constraint_violations_pullback
end
function ChainRulesCore.rrule(::typeof(BatchNLPKernels.bound_violations!), bm::BatchModel, X)
    Vb = BatchNLPKernels.bound_violations!(bm, X)

    function bound_violations_pullback(Ȳ)
        Ȳ = ChainRulesCore.unthunk(Ȳ)

        # violation(x, s) = max(s.l - x, x - s.u, 0)
        # ∂violation/∂x = -1 if x < s.l, +1 if x > s.u, 0 otherwise

        X̄ = if isempty(bm.viols_vars)
            zeros(eltype(X), size(X))
        else
            lower_viols = X .< bm.viols_vars.l
            upper_viols = X .> bm.viols_vars.u
            lower_viols .* (-Ȳ) .+ upper_viols .* Ȳ
        end

        return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), X̄
    end

    return Vb, bound_violations_pullback
end

end # module BNKChainRulesCore