trainables

Author: CarloLucibello

src/destructure.jl

@@ -73,7 +73,7 @@ function _flatten(x)
     o
   end
   isempty(arrays) && return Bool[], off, 0
-  reduce(vcat, arrays), off, len[]
+  return reduce(vcat, arrays), off, len[]
 end
 
 struct _TrainableStructWalk <: AbstractWalk end
@@ -174,3 +174,4 @@ function ChainRulesCore.rrule(::typeof(_maybewarn))
   @warn "second derivatives of destructure may not work yet, sorry!" maxlog=3
   nothing, _ -> (NoT,)
 end
+

src/interface.jl

@@ -167,6 +167,7 @@ and `trainable(x)` must contain a subset of these.
 """
 trainable(x) = functor(x)[1]
 
+# like trainable(x), but also tries to output non-trainable children giving value nothing
 _trainable(x) = _trainable(functor(x)[1], trainable(x))
 _trainable(ch::NamedTuple, tr::NamedTuple) = merge(map(_ -> nothing, ch), tr)
 _trainable(ch::Tuple{Vararg{Any,N}}, tr::Tuple{Vararg{Any,N}}) where N = tr

src/trainables.jl

@@ -0,0 +1,35 @@
+
+
+function trainables1(x)
+    isnumeric(x) && return [x]
+    arrays = AbstractArray[]
+    fmap(x; exclude = isnumeric, walk = _TrainableStructWalk()) do y
+        push!(arrays, y)
+        return y
+    end
+    return arrays
+end
+
+############
+
+using Functors: AbstractWalk, _map, _values, execute, ExcludeWalk
+
+struct TrainableWalk2 <: AbstractWalk end
+
+function (walk::TrainableWalk2)(recurse, x, ys...)
+    x_children = _values(Optimisers.trainable(x))
+    ys_children = map(Optimisers.trainable, ys)
+    res = _map(recurse, x_children, ys_children...)
+    @show _values(res)
+    return _values(res)
+end
+
+function trainables2(x)
+    exclude(x) = Optimisers.isnumeric(x) && Functors.isleaf(x)
+    return execute(ExcludeWalk(TrainableWalk2(), x -> x, exclude), x)
+end
+
+using Flux
+
+m = Chain(Dense(2 => 3, relu), BatchNorm(3), Dense(3 => 2))
+trainables2(m)