feat: Added several tests for correctness.

src/dense.jl

@@ -1,30 +1,58 @@
+"""
+    NIG(in => out, σ=NNlib.softplut; bias=true, init=Flux.glorot_uniform)
+    NIG(W::AbstractMatrix, [bias, σ])
 
+Create a fully connected layer which implements the NormalInverseGamma Evidential distribution
+whose forward pass is simply given by:
+
+    y = W * x .+ bias
+
+The input `x` should be a vector of length `in`, or batch of vectors represented
+as an `in × N` matrix, or any array with `size(x,1) == in`.
+The out `y` will be a vector  of length `out*4`, or a batch with
+`size(y) == (out*4, size(x)[2:end]...)`
+The output will have applied the function `σ(y)` to each row/element of `y` except the first `out` ones.
+Keyword `bias=false` will switch off trainable bias for the layer.
+The initialisation of the weight matrix is `W = init(out*4, in)`, calling the function
+given to keyword `init`, with default [`glorot_uniform`](@doc Flux.glorot_uniform).
+The weight matrix and/or the bias vector (of length `out`) may also be provided explicitly.
+Remember that in this case the number of rows in the weight matrix `W` MUST be a multiple of 4.
+The same holds true for the `bias` vector.
+
+# Arguments:
+- `(in, out)`: number of input and output neurons
+- `σ`: The function to use to secure positive only outputs which defaults to the softplus function.
+- `init`: The function to use to initialise the weight matrix.
+- `bias`: Whether to include a trainable bias vector.
+"""
 struct NIG{F,M<:AbstractMatrix,B}
         W::M
         b::B
         σ::F
-        function NIG(W::M, b = true, σ::F = identity) where {M<:AbstractMatrix,F}
+        function NIG(W::M, b = true, σ::F = NNlib.softplus) where {M<:AbstractMatrix,F}
                 b = Flux.create_bias(W, b, size(W, 1))
                 return new{F,M,typeof(b)}(W, b, σ)
         end
 end
 
-function NIG((in, out)::Pair{<:Integer,<:Integer}, σ = identity;
+function NIG((in, out)::Pair{<:Integer,<:Integer}, σ = NNlib.softplus;
         init = Flux.glorot_uniform, bias = true)
         NIG(init(out * 4, in), bias, σ)
 end
 
 Flux.@functor NIG
 
 function (a::NIG)(x::AbstractVecOrMat)
-        o = a.σ.(a.W * x .+ a.b)
-        μ = @view o[1:4, :]
-        ν = @view o[5:8, :]
-        ν .= NNlib.softplus.(ν)
-        α = @view o[9:12, :]
-        α .= NNlib.softplus.(α) .+ 1
-        β = @view o[13:16, :]
-        β .= NNlib.softplus.(β)
-        return o
+	nout = Int(size(a.W, 1) / 4)
+        o = a.W * x .+ a.b
+        γ = o[1:nout, :]
+	ν = o[(nout+1):(nout*2), :]
+        ν = a.σ.(ν)
+	α = o[(nout*2+1):(nout*3), :]
+        α = a.σ.(α) .+ 1
+	β = o[(nout*3+1):(nout*4), :]
+        β = a.σ.(β)
+        return vcat(γ, ν, α, β)
 end
+
 (a::NIG)(x::AbstractArray) = reshape(a(reshape(x, size(x, 1), :)), :, size(x)[2:end]...)

test/runtests.jl

@@ -1,12 +1,63 @@
 using EvidentialFlux
+using Flux
 using Test
 
 @testset "EvidentialFlux.jl" begin
-    # Write your tests here.
-    m = NIG(3 => 4)
-    x = randn(Float32, 3, 10)
-    y = m(x)
-    @test size(y) == (16, 10)
-    @test y[5:16, :] == abs.(y[5:16, :])
-    #@test y[9:12, :] .> 1
+        # Creating a model and a forward pass
+
+        ninp, nout = 3, 5
+        m = NIG(ninp => nout)
+        x = randn(Float32, 3, 10)
+        ŷ = m(x)
+        @test size(ŷ) == (20, 10)
+        # The ν, α, and β all have to be positive
+        @test ŷ[6:20, :] == abs.(ŷ[6:20, :])
+        # The α all have to be ≥ 1
+        @test all(≥(1), ŷ[11:15, :])
+
+        # Testing backward pass
+        oldW = similar(m.W)
+        oldW .= m.W
+        loss(y, ŷ) = sum(abs, y - ŷ)
+        pars = Flux.params(m)
+        y = randn(Float32, nout, 10) # Target (fake)
+        grads = Flux.gradient(pars) do
+                ŷ = m(x)
+                γ = ŷ[1:nout, :]
+                loss(y, γ)
+        end
+        # Test that we can update the weights based on gradients
+        opt = Descent(0.1)
+        Flux.Optimise.update!(opt, pars, grads)
+        @test m.W != oldW
+
+        # Testing convergence
+        ninp, nout = 3, 1
+        m = NIG(ninp => nout)
+        x = Float32.(collect(1:0.1:10))
+        x = cat(x', x' .- 10, x' .+ 5, dims = 1)
+        # y = 1 * sin.(x[1, :]) .- 3 * sin.(x[2, :]) .+ 2 * cos.(x[3, :]) .+ randn(Float32, 91)
+        y = 1 * x[1, :] .- 3 * x[2, :] .+ 2 * x[3, :] .+ randn(Float32, 91)
+        #scatterplot(x[1, :], y, width = 90, height = 30)
+        #loss(y, ŷ) = sum(abs, y - ŷ)
+        pars = Flux.params(m)
+        opt = ADAMW(0.005)
+	trnlosses = zeros(Float32, 1000)
+        for i in 1:1000
+		local trnloss
+                grads = Flux.gradient(pars) do
+                        ŷ = m(x)
+                        γ = ŷ[1, :]
+                        trnloss = loss(y, γ)
+                end
+		trnlosses[i] = trnloss
+                # Test that we can update the weights based on gradients
+                Flux.Optimise.update!(opt, pars, grads)
+		#if i % 100 == 0 
+		#	println("Epoch $i, Loss: $trnloss")
+		#end
+        end
+	#scatterplot(1:1000, trnlosses)
+	@test trnlosses[10] > trnlosses[100] > trnlosses[300]
+
 end