WIP: Make optimize work on structs

src/optimise/train.jl

@@ -17,17 +17,36 @@ Update the array `x` according to `x .-= x̄`.
 """
 function update!(x::AbstractArray, x̄)
   x .-= x̄
+  return
 end
 
-function update!(opt, x, x̄)
-  x .-= apply!(opt, x, x̄)
+# skip, if gradient is nothing
+update!(x::AbstractArray, x̄::Nothing) = nothing
+update!(opt, x::AbstractArray, x̄::Nothing) = nothing
+update!(opt, m::M, ∇m::Nothing) where M = nothing
+
+function update!(opt, x::AbstractArray, x̄)
+  x .-= apply!(opt, x::AbstractArray, x̄)
+  return
+end
+
+# NOTE: since there won't be real loop in a struct
+#       we could always flatten it, which is a bit
+#       faster.
+@generated function update!(opt, m::M, ∇m) where M
+  body = Expr(:block)
+  for each in fieldnames(M)
+    each = QuoteNode(each)
+    push!(body.args, :(update!(opt, getfield(m, $each), getfield(∇m, $each))))
+  end
+  return body
 end
 
 function update!(opt, xs::Params, gs)
   for x in xs
-    gs[x] == nothing && continue
     update!(opt, x, gs[x])
   end
+  return
 end
 
 # Callback niceties