add fully connect nn constructor

Author: zuhengxu

src/flows/utils.jl

@@ -2,6 +2,8 @@ using Bijectors: transformed
 using Flux
 
 """
+    mlp3(input_dim::Int, hidden_dims::Int, output_dim::Int; activation=Flux.leakyrelu)
+
 A simple wrapper for a 3 layer dense MLP
 """
 function mlp3(input_dim::Int, hidden_dims::Int, output_dim::Int; activation=Flux.leakyrelu)
@@ -12,7 +14,52 @@ function mlp3(input_dim::Int, hidden_dims::Int, output_dim::Int; activation=Flux
     )
 end
 
+"""
+    fnn(
+        input_dim::Int,
+        hidden_dims::AbstractVector{<:Int},
+        output_dim::Int;
+        inlayer_activation=Flux.leakyrelu,
+        output_activation=Flux.tanh,
+    )
+
+Create a fully connected neural network (FNN).
+
+# Arguments
+- `input_dim::Int`: The dimension of the input layer.
+- `hidden_dims::AbstractVector{<:Int}`: A vector of integers specifying the dimensions of the hidden layers.
+- `output_dim::Int`: The dimension of the output layer.
+- `inlayer_activation`: The activation function for the hidden layers. Defaults to `Flux.leakyrelu`.
+- `output_activation`: The activation function for the output layer. Defaults to `Flux.tanh`.
+
+# Returns
+- A `Flux.Chain` representing the FNN.
+"""
+function fnn(
+    input_dim::Int,
+    hidden_dims::AbstractVector{<:Int},
+    output_dim::Int;
+    inlayer_activation=Flux.leakyrelu,
+    output_activation=Flux.tanh,
+)
+    # Create a chain of dense layers
+    # First layer
+    layers = Any[Flux.Dense(input_dim, hidden_dims[1], inlayer_activation)]
+
+    # Hidden layers
+    for i in 1:(length(hidden_dims) - 1)
+        push!(
+            layers,
+            Flux.Dense(hidden_dims[i], hidden_dims[i + 1], inlayer_activation),
+        )
+    end
+
+    # Output layer
+    push!(layers, Flux.Dense(hidden_dims[end], output_dim, output_activation))
+    return Chain(layers...)
+end
+
 function create_flow(Ls, q₀)
     ts =  reduce(∘, Ls)
     return transformed(q₀, ts)
-end
+end