|
| 1 | +using Flux |
| 2 | +using Bijectors |
| 3 | +using Bijectors: partition, combine, PartitionMask |
| 4 | + |
| 5 | +using Random, Distributions, LinearAlgebra |
| 6 | +using Functors |
| 7 | +using Optimisers, ADTypes |
| 8 | +using Mooncake, Zygote, Enzyme, ADTypes |
| 9 | +import NormalizingFlows as NF |
| 10 | + |
| 11 | +import DifferentiationInterface as DI |
| 12 | + |
| 13 | + |
| 14 | +pt = Float64 |
| 15 | +inputdim = 4 |
| 16 | +outputdim = 3 |
| 17 | + |
| 18 | +x = randn(pt, inputdim) |
| 19 | + |
| 20 | +bs = 64 |
| 21 | +xs = randn(pt, inputdim, 64) |
| 22 | + |
| 23 | +# compose two fully connected networks |
| 24 | +m1 = NF.fnn(inputdim, [16, 16], outputdim; output_activation=nothing, paramtype=pt) |
| 25 | +m2 = NF.fnn(outputdim, [16, 16], inputdim; output_activation=Flux.tanh, paramtype=pt) |
| 26 | +mm = reduce(∘, (m2, m1)) |
| 27 | +psm, stm = Optimisers.destructure(mm) |
| 28 | + |
| 29 | +function lsm(ps, st, x) |
| 30 | + model = st(ps) |
| 31 | + y = model(x) |
| 32 | + return sum(y) # just a dummy loss |
| 33 | +end |
| 34 | + |
| 35 | +adtype = ADTypes.AutoMooncake(; config = Mooncake.Config()) |
| 36 | + |
| 37 | +val, grad = DI.value_and_gradient( |
| 38 | + lsm, adtype, |
| 39 | + psm, DI.Cache(stm), DI.Constant(xs) |
| 40 | +) |
| 41 | + |
| 42 | + |
| 43 | +acl = NF.AffineCoupling( inputdim, [16, 16], 1:2:inputdim, pt) |
| 44 | +psacl,stacl = Optimisers.destructure(acl) |
| 45 | + |
| 46 | +function loss(ps, st, x) |
| 47 | + model = st(ps) |
| 48 | + y = model(x) |
| 49 | + return sum(y) # just a dummy loss |
| 50 | +end |
| 51 | + |
| 52 | +val, grad = DI.value_and_gradient( |
| 53 | + loss, |
| 54 | + ADTypes.AutoEnzyme(; |
| 55 | + mode=Enzyme.set_runtime_activity(Enzyme.Reverse), |
| 56 | + function_annotation=Enzyme.Const, |
| 57 | + ), |
| 58 | + psacl, DI.Cache(stacl), DI.Constant(x) |
| 59 | +) |
| 60 | + |
| 61 | +# val, grad = DI.value_and_gradient( |
| 62 | +# loss, |
| 63 | +# ADTypes.AutoMooncake(; config = Mooncake.Config()), |
| 64 | +# psacl, DI.Cache(stacl), DI.Constant(x) |
| 65 | +# ) |
| 66 | + |
| 67 | +function loss_acl_manual(ps, st, x) |
| 68 | + acl = st(ps) |
| 69 | + s_net = acl.s |
| 70 | + t_net = acl.t |
| 71 | + mask = acl.mask |
| 72 | + x₁, x₂, x₃ = partition(mask, x) |
| 73 | + y₁ = exp.(s_net(x₂)) .* x₁ .+ t_net(x₂) |
| 74 | + y = combine(mask, y₁, x₂, x₃) |
| 75 | + # println("y = ", y) |
| 76 | + return sum(y) |
| 77 | +end |
| 78 | + |
| 79 | +val, grad = DI.value_and_gradient( |
| 80 | + loss_acl_manual, |
| 81 | + # ADTypes.AutoMooncake(; config = Mooncake.Config()), |
| 82 | + # ADTypes.AutoEnzyme(; |
| 83 | + # mode=Enzyme.set_runtime_activity(Enzyme.Reverse), |
| 84 | + # function_annotation=Enzyme.Const, |
| 85 | + # ), |
| 86 | + psacl, DI.Cache(stacl), DI.Constant(x) |
| 87 | +) |
| 88 | + |
| 89 | + |
| 90 | + |
| 91 | +function mlp3( |
| 92 | + input_dim::Int, |
| 93 | + hidden_dims::Int, |
| 94 | + output_dim::Int; |
| 95 | + activation=Flux.leakyrelu, |
| 96 | + paramtype::Type{T} = Float64 |
| 97 | +) where {T<:AbstractFloat} |
| 98 | + m = Chain( |
| 99 | + Flux.Dense(input_dim, hidden_dims, activation), |
| 100 | + Flux.Dense(hidden_dims, hidden_dims, activation), |
| 101 | + Flux.Dense(hidden_dims, output_dim), |
| 102 | + ) |
| 103 | + return Flux._paramtype(paramtype, m) |
| 104 | +end |
| 105 | + |
| 106 | +function ls_msk(ps, st, x, mask) |
| 107 | + t_net = st(ps) |
| 108 | + x₁, x₂, x₃ = partition(mask, x) |
| 109 | + y₁ = x₁ .+ t_net(x₂) |
| 110 | + y = combine(mask, y₁, x₂, x₃) |
| 111 | + # println("y = ", y) |
| 112 | + return sum(abs2, y) |
| 113 | +end |
| 114 | + |
| 115 | +inputdim = 4 |
| 116 | +mask_idx = 1:2:inputdim |
| 117 | +mask = PartitionMask(inputdim, mask_idx) |
| 118 | +cdim = length(mask_idx) |
| 119 | + |
| 120 | +x = randn(inputdim) |
| 121 | + |
| 122 | +t_net = mlp3(cdim, 16, cdim; paramtype = Float64) |
| 123 | +ps, st = Optimisers.destructure(t_net) |
| 124 | + |
| 125 | +ls_msk(ps, st, x, mask) # 3.0167880799441793 |
| 126 | + |
| 127 | +val, grad = DI.value_and_gradient( |
| 128 | + ls_msk, |
| 129 | + ADTypes.AutoMooncake(; config = Mooncake.Config()), |
| 130 | + ps, DI.Cache(st), DI.Constant(x), DI.Constant(mask) |
| 131 | +) |
| 132 | + |
| 133 | + |
| 134 | +struct ACL |
| 135 | + mask::Bijectors.PartitionMask |
| 136 | + t::Flux.Chain |
| 137 | +end |
| 138 | +@functor ACL (t, ) |
| 139 | + |
| 140 | +acl = ACL(mask, t_net) |
| 141 | +psacl, stacl = Optimisers.destructure(acl) |
| 142 | + |
| 143 | +function loss_acl(ps, st, x) |
| 144 | + acl = st(ps) |
| 145 | + t_net = acl.t |
| 146 | + mask = acl.mask |
| 147 | + x₁, x₂, x₃ = partition(mask, x) |
| 148 | + y₁ = x₁ .+ t_net(x₂) |
| 149 | + y = combine(mask, y₁, x₂, x₃) |
| 150 | + return sum(abs2, y) |
| 151 | +end |
| 152 | +loss_acl(psacl, stacl, x) # 3.0167880799441793 |
| 153 | + |
| 154 | +val, grad = DI.value_and_gradient( |
| 155 | + loss_acl, |
| 156 | + ADTypes.AutoEnzyme(; |
| 157 | + mode=Enzyme.set_runtime_activity(Enzyme.Reverse), |
| 158 | + function_annotation=Enzyme.Const, |
| 159 | + ), |
| 160 | + psacl, DI.Cache(stacl), DI.Constant(x) |
| 161 | +) |
| 162 | + |
| 163 | +val, grad = DI.value_and_gradient( |
| 164 | + loss_acl, |
| 165 | + ADTypes.AutoMooncake(; config = Mooncake.Config()), |
| 166 | + psacl, DI.Cache(stacl), DI.Constant(x) |
| 167 | +) |
0 commit comments