# Lifting a Julia function into a Flux "layer"

**URL:** <https://discourse.julialang.org/t/lifting-a-julia-function-into-a-flux-layer/13712>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [August 19, 2018, 3:22pm UTC](https://discourse.julialang.org/t/lifting-a-julia-function-into-a-flux-layer/13712 "2018-08-19T15:22:23Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![cscherrer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cscherrer/32/7631_2.png) [@cscherrer](https://discourse.julialang.org/u/cscherrer)\
**Post date:** [August 19, 2018, 3:22pm UTC](https://discourse.julialang.org/t/lifting-a-julia-function-into-a-flux-layer/13712/1 "2018-08-19T15:22:23Z")

</div>

Flux confused me for a long time. I got the first example in the docs to work:

```julia
using Flux.Tracker

W = param(rand(2, 5))
b = param(rand(2))

predict(x) = W*x .+ b
loss(x, y) = sum((predict(x) .- y).^2)

x, y = rand(5), rand(2) # Dummy data
l = loss(x, y) # ~ 3

params = Params([W, b])
grads = Tracker.gradient(() -> loss(x, y), params)

```

But other examples use a `params` function. This seems to be required to use optimizers like SGD and ADAM.

I spent a few hours trying to use ADAM to optimize a simple example, with no luck. Flux’s optimization seems to require a special “layer”, which is not just a function, but requires some kind of wrapper.

Is there a way to “lift” an arbitrary function so Flux will treat it as a layer?

---

<div class="post-metadata">

**Author:** ![baggepinnen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/baggepinnen/32/693_2.png) [@baggepinnen](https://discourse.julialang.org/u/baggepinnen)\
**Post date:** [August 19, 2018, 4:20pm UTC](https://discourse.julialang.org/t/lifting-a-julia-function-into-a-flux-layer/13712/2 "2018-08-19T16:20:56Z")

</div>

You need to call the function with `Param`s so that flux can track the gradients. `params` only returns the `Param`s inside a struct.

---

<div class="post-metadata">

**Author:** ![cscherrer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cscherrer/32/7631_2.png) [@cscherrer](https://discourse.julialang.org/u/cscherrer)\
**Post date:** [August 19, 2018, 6:24pm UTC](https://discourse.julialang.org/t/lifting-a-julia-function-into-a-flux-layer/13712/3 "2018-08-19T18:24:24Z")

</div>

Ok, so the docs have this example:

```julia
opt = SGD([W, b], 0.1) # Gradient descent with learning rate 0.1

opt() # Carry out the update, modifying `W` and `b`.

```

So let’s try it for the case above:

```julia
using Flux
using Flux.Tracker
using Flux.Optimise

W = param(rand(2, 5))
b = param(rand(2))

predict(x) = W*x .+ b
loss(x, y) = sum((predict(x) .- y).^2)

x, y = rand(5), rand(2) # Dummy data

par = Params([W, b])
grads = Tracker.gradient(() -> loss(x, y), par)

opt = SGD(par, 0.1)
opt()

```

This returns

```julia
MethodError: no method matching length(::Params)
Closest candidates are:
  length(!Matched::Core.SimpleVector) at essentials.jl:576
  length(!Matched::Base.MethodList) at reflection.jl:728
  length(!Matched::Core.MethodTable) at reflection.jl:802
  ...

```

But it doesn’t make sense, anyway - the call uses parameters, but has no reference to the loss function.

It seems that for proper Flux model, `[W,b]` or `params(m)` must store the loss function under the hood. But it’s not clear how to connect this for a hand-rolled loss function.

---

<div class="post-metadata">

**Author:** ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)\
**Post date:** [August 19, 2018, 9:50pm UTC](https://discourse.julialang.org/t/lifting-a-julia-function-into-a-flux-layer/13712/4 "2018-08-19T21:50:20Z")

</div>

> [@cscherrer](#):
>
> It seems that for proper Flux model, `[W,b]` or `params(m)` must store the loss function under the hood.

Not necessary. In many ML libs nowadays optimizers are as simple as a single function that takes 2 arguments - parameter and its update - and applies one to another according to optimizer-specific rules (e.g. for GD it’s just x := x - \alpha \Delta). Flux’s `Param` struct already [contains parameter value and its delta](https://github.com/FluxML/Flux.jl/blob/4045c322d520b2bd63efe0a46b072d5f03eeccf4/src/optimise/Optimise.jl#L7-L10), so optimizer doesn’t care about loss function or whatever.

---

<div class="post-metadata">

**Author:** ![cscherrer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cscherrer/32/7631_2.png) [@cscherrer](https://discourse.julialang.org/u/cscherrer)\
**Post date:** [August 19, 2018, 11:07pm UTC](https://discourse.julialang.org/t/lifting-a-julia-function-into-a-flux-layer/13712/5 "2018-08-19T23:07:51Z")

</div>

> [@dfdx](#):
>
> In many ML libs nowadays optimizers are as simple as a single function that takes 2 arguments - parameter and its update

Now I’m really confused. I can imagine precomputing the gradient and having an optimization algorithm maintain state including momentum, learning rate, etc. But this sounds like the not only the gradient, but also the update is computed before passing to the optimizer.

So, when you call `params(m)`, it computes what update to apply, independent of any context?

I would really appreciate if anyone can post an “easy example done the hard way”, maybe something like a simple linear or logistic regression using the SGD function from Flux.

---

<div class="post-metadata">

**Author:** ![baggepinnen](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/baggepinnen/32/693_2.png) [@baggepinnen](https://discourse.julialang.org/u/baggepinnen)\
**Post date:** [August 20, 2018, 7:06am UTC](https://discourse.julialang.org/t/lifting-a-julia-function-into-a-flux-layer/13712/6 "2018-08-20T07:06:04Z")

</div>

> [@cscherrer](#):
>
> So, when you call `params(m)` , it computes what update to apply, independent of any context?

No, `params(m)` only returns the parameters stored inside struct `m`.  
When `y = f(x::TrackedArray)` is called, the output `f(x)` is returned as a tracked value `y`, and the gradient of `f` w.r.t. `x` is stored inside `x.grad` (technically, I think you have to call `back(y)` for the gradients to be computed). When `opt` is called, it takes `x.grad` and modifies `x` according to

1. `x.grad`
2. ´The specific learning rule of the optimizer, e.g., ADAM, SGD etc.

So the whole thing hinges on the fact that you do your calculations using `TrackedArray/TrackedReal` so that Flux can keep track of the gradients. The optimizers store a pointer to the parameters that they are supposed to update, and the gradients are stored inside these `TrackedArray`s.

---

<div class="post-metadata">

**Author:** ![cscherrer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/cscherrer/32/7631_2.png) [@cscherrer](https://discourse.julialang.org/u/cscherrer)\
**Post date:** [August 20, 2018, 3:36pm UTC](https://discourse.julialang.org/t/lifting-a-julia-function-into-a-flux-layer/13712/7 "2018-08-20T15:36:34Z")

</div>

> [@baggepinnen](#):
>
> No, `params(m)` only returns the parameters stored inside struct `m` .

I’m not sure any more what you mean by “parameters”, since it’s so overloaded.

- `Param` contains a value (of arbitrary type) and a “\Delta” of the same type, which is somehow related to how the value is to be updated, and so maybe to its gradient.
- `Params` is another struct, containing only a value called `params`.
- But `params` is also a function that can be called on a model.
- `params` is also the name of the argument passed to optimizers like `SGD`.

It seems to me that `param` says “this thing should be considered a parameter” and returns a `TrackedArray`, in order to track gradients.

I’ve gone through the docs some more, and I think I’m making progress. There’s this example:

```julia
struct Affine
  W
  b
end

Affine(in::Integer, out::Integer) =
  Affine(param(randn(out, in)), param(randn(out)))

# Overload call, so the object can be used as a function
(m::Affine)(x) = m.W * x .+ m.b

a = Affine(10, 5)

a(rand(10)) # => 5-element vector

```

At the bottom it mentions [`treelike`](http://fluxml.ai/Flux.jl/latest/models/basics.html#Layer-helpers-1), which seems to be the important piece I was missing. I guess it’s maybe better to think “`treelike!`”, because its side effect is the whole point. Once you call this, you can call `train!` and it seems to actually work.

I’ll try to build logistic regression in this way and see how it goes. Thank you @dfdx and @baggepinnen for your help!

---

<div class="post-metadata">

**Author:** ![robsmith11](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/robsmith11/32/29641_2.png) [@robsmith11](https://discourse.julialang.org/u/robsmith11)\
**Post date:** [May 29, 2019, 4:11pm UTC](https://discourse.julialang.org/t/lifting-a-julia-function-into-a-flux-layer/13712/8 "2019-05-29T16:11:11Z")

</div>

I just started playing with Flux and am confused by the meaning/purpose of `Params` as well. The intro docs are not clear at all on this.