# Laplacian of shallow Flux neural network with complex weights

**URL:** https://discourse.julialang.org/t/laplacian-of-shallow-flux-neural-network-with-complex-weights/72600
**Category:** General Usage
**Created:** [December 5, 2021, 4:14am UTC](https://discourse.julialang.org/t/laplacian-of-shallow-flux-neural-network-with-complex-weights/72600 "2021-12-05T04:14:06Z")
**Posts on this page:** 1
**Page:** 1

<div class="post-metadata">

### Author: ![Matematija](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/matematija/32/21820_2.png) [@Matematija](https://discourse.julialang.org/u/Matematija)
#### Post date: [December 5, 2021, 4:14am UTC](https://discourse.julialang.org/t/laplacian-of-shallow-flux-neural-network-with-complex-weights/72600/1 "2021-12-05T04:14:06Z")

</div>

Hi all,

I am trying to use Julia’s AD ecosystem to calculate what the title says - a Laplacian (trace of Hessian matrix) of a shallow (1 layer) Flux neural net.

There are two twists to the story:

- I need all of the weights to be complex because of domain-specific reasons. My inputs are still real though.
- The “activation” function actually comes from `SpecialFunctions.jl`, which seems to block AD with `ForwardDiff.jl` (I raised an [issue on their GitHub](https://github.com/JuliaDiff/ForwardDiff.jl/issues/567))

Here is a simplified version of the model with auxiliary functions:

```julia
using Flux
using SpecialFunctions
using Zygote
using ForwardDiff

log_i0(z) = log(besselix(0,z)) .+ abs(real(z))

complex_init(dims...) = convert.(ComplexF64, Flux.glorot_normal(dims)) .* exp.(2π*im .*rand(Float64, dims...))

model = Chain(
    complex,
    Conv((3,3), 1=>4, log_i0; init=complex_init),
    sum
)

inputs = randn(8, 8 ,1, 1)
model(inputs)

```

First derivatives w.r.t inputs I can easily get through `Zygote`, no issues there:

```julia
Zygote.gradient(real∘model, complex(inputs)) |> first |> conj

```

My question is - can I get trace of Hessian matrix (Laplacian) of such a model? Are there tricks or different packages I am not aware of?

Something like `Zygote.diaghessian` looks perfect but uses `ForwardDiff` under the hood which has a problem with `log_i0` and, even if I replace it with something less exotic, throws other errors.