rename MLP_3layer to mlp3 for convenience

Author: zuhengxu

example/HamiltonianVI/model.jl

@@ -1,4 +1,22 @@
-using Random, Distributions, LinearAlgebra, Bijectors
+using Random, Distributions, LinearAlgebra
+using Bijectors: transformed
+using Flux
+
+"""
+A simple wrapper for a 3 layer dense MLP
+"""
+function mlp3(input_dim::Int, hidden_dims::Int, output_dim::Int; activation=Flux.leakyrelu)
+    return Chain(
+        Flux.Dense(input_dim, hidden_dims, activation),
+        Flux.Dense(hidden_dims, hidden_dims, activation),
+        Flux.Dense(hidden_dims, output_dim),
+    )
+end
+
+function create_flow(Ls, q₀)
+    ts =  reduce(∘, Ls)
+    return transformed(q₀, ts)
+end
 
 function compare_trained_and_untrained_flow(
     flow_trained::Bijectors.MultivariateTransformed,
@@ -79,10 +97,6 @@ end
 #     return p
 # end
 
-function create_flow(Ls, q₀)
-    ts =  reduce(∘, Ls)
-    return transformed(q₀, ts)
-end
 
 function visualize(p::Bijectors.MultivariateTransformed, samples=rand(p, 1000))
     xrange = range(minimum(samples[1, :]) - 1, maximum(samples[1, :]) + 1; length=100)

example/HamiltonianVI/momentum.jl

@@ -32,8 +32,8 @@ function AffineCoupling(
     mask_idx::AbstractVector, # index of dimensione that one wants to apply transformations on
 )
     cdims = length(mask_idx) # dimension of parts used to construct coupling law
-    s = MLP_3layer(cdims, hdims, cdims)
-    t = MLP_3layer(cdims, hdims, cdims)
+    s = mlp3(cdims, hdims, cdims)
+    t = mlp3(cdims, hdims, cdims)
     mask = PartitionMask(dim, mask_idx)
     return AffineCoupling(dim, mask, s, t)
 end
@@ -99,8 +99,8 @@ end
 #     mask_idx::AbstractVector, # index of dimensione that one wants to apply transformations on
 # )
 #     cdims = length(mask_idx) # dimension of parts used to construct coupling law
-#     s = MLP_3layer(cdims, hdims, cdims)
-#     t = MLP_3layer(cdims, hdims, cdims)
+#     s = mlp3(cdims, hdims, cdims)
+#     t = mlp3(cdims, hdims, cdims)
 #     mask = PartitionMask(dim, mask_idx)
 #     return AffineCoupling(dim, mask, s, t)
 # end

example/common.jl

@@ -39,7 +39,7 @@ function NeuralSplineLayer(
 ) where {T1<:Int,T2<:Real}
     num_of_transformed_dims = length(mask_idx)
     input_dims = dim - num_of_transformed_dims
-    nn = [MLP_3layer(input_dims, hdims, 3K - 1) for _ in 1:num_of_transformed_dims]
+    nn = [mlp3(input_dims, hdims, 3K - 1) for _ in 1:num_of_transformed_dims]
     mask = Bijectors.PartitionMask(dim, mask_idx)
     return NeuralSplineLayer(dim, K, nn, B, mask)
 end