# AD Troubles in Flux and Unusual Loss

**URL:** <https://discourse.julialang.org/t/ad-troubles-in-flux-and-unusual-loss/122995>\
**Category:** Machine Learning\
**Tags:** flux, zygote, fluxtraining\
**Created:** [November 23, 2024, 2:12pm UTC](https://discourse.julialang.org/t/ad-troubles-in-flux-and-unusual-loss/122995 "2024-11-23T14:12:30Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![gideonsimpson](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gideonsimpson/32/1928_2.png) [@gideonsimpson](https://discourse.julialang.org/u/gideonsimpson)\
**Post date:** [November 23, 2024, 2:12pm UTC](https://discourse.julialang.org/t/ad-troubles-in-flux-and-unusual-loss/122995/1 "2024-11-23T14:12:30Z")

</div>

I’m struggling with the following issue. I want to train on a somewhat nonstandard loss function involving the gradient of the target. If everything _worked_, I would. use the following code:

```julia
using Flux
using Random
using Printf
using LinearAlgebra
using Statistics

d = 2;
h = 1.0f0;
n_x = 10;

f_true(x) = sin(x[1]) * cos(x[2]); # arbitrary, just want a 2=>1 example function

# cook up data
Random.seed!(500);

x_train = [randn(Float32, d) for _ in 1:n_x];
y_train = Float32.(f_true.(x_train));

# convert to matrix
X_train = hcat(x_train...);
Y_train = hcat(y_train...);

W_train = [X_train[:, j] .- X_train for j in 1:n_x];
training_data = [(W_train, y_train)];

inner_layer = Dense(d => d, bias=false);
# what I would naively implement, to create a column vector of the gaussian of the column-wise norms
function problem_layer1(w)
    return [exp(-norm(w[:, j])^2 / (2h)) for j in 1:n_x]
end
outer_layer = Dense(n_x => 1, bias=false);
f1 = Chain(inner_layer, problem_layer1, outer_layer);

# nonstandard loss
function grad_loss(f, w, y)
    ∇g(w) = Flux.gradient(ϕ_ -> Chain(f.layers[2:end])(ϕ_)[1][1], f.layers[1](w))[1]
    return mean([norm(∇g(w_))^2 for w_ in w])
end
# now try training
opt_f_state1 = Flux.setup(Flux.Adam(0.01), f1)
Flux.train!(grad_loss, f1, training_data, opt_f_state1)

```

and this produces the error:

```julia
Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_eqtable_get), Any, svec(Any, Any, Any), 0, :(:ccall), %5, %3, %4)).

```

Next, I tried something slightly different, which did work:

```julia
# this is not the same function, but it has a similar flavor
function problem_layer2(w)
    return sum(abs2, w, dims=1)'
end
g2 = Chain(inner_layer, problem_layer2, outer_layer);
# try training
opt_g_state2 = Flux.setup(Flux.Adam(0.01), g2)
Flux.train!(grad_loss, g2, training_data, opt_g_state2)

```

and this executes without error (which is not to say it’s _correct_).

Next, I tried:

```julia
f2 = Chain(inner_layer, problem_layer2, norm2 -> exp.(-norm2 ./ (2h)), outer_layer);

opt_f_state2 = Flux.setup(Flux.Adam(0.01), f2)
Flux.train!(grad_loss, f2, training_data, opt_f_state2)

```

and I got

```julia
Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_eqtable_get), Any, svec(Any, Any, Any), 0, :(:ccall), %5, %3, %4)).

```

I even tried the simpler problem

```julia
h2 = Chain(inner_layer, problem_layer2, norm2 -> norm2 ./ (-2h), outer_layer);

```

which just does a scaling, and this also gave me a `foreigncall` error.

Thanks for any suggestions on how to sort this out.

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [November 24, 2024, 12:09am UTC](https://discourse.julialang.org/t/ad-troubles-in-flux-and-unusual-loss/122995/2 "2024-11-24T00:09:08Z")

</div>

Flux (Zygote) doesn’t deal with nested gradients nicely.

You need to either use Enzyme for the autodiff or if you really want to use Zygote, then call it via Lux ([Nested Automatic Differentiation | Lux.jl Docs](https://lux.csail.mit.edu/stable/manual/nested_autodiff))