Bayesian neural networks in Flux

I am trying to implement a Bayesian neural network in Flux. I have custom layers which return two outputs. One is the output of the layer and another is a conditional probability. I use these probabilities to define a loss function using the negative log-likelihood. My question is that how can I define a custom layer with two outputs instead of one?

1 Like

just return both: return out1, out2
If the layer is stacked into a Chain, make sure the following layer can work with an input tuple of length 2

So is this code true? All the layers in the network are of the same type.


> struct BayesianLayer
>        model
> end
> function (l::BayesianLayer)(x)
>       len = length(x)
>       if len == 1
>                out1 = l.model(x) 
>                out2 = somefunction(out1)
>       elseif  len == 2
>                out1 = l.model(x[1])
>                out2 = somefunction(out1)*x[2]
>       end
>       return out1, out2
> end