Add lazy broadcasting (#31)

* add lazy broadcasting

* magic copy keyword

* Revert "magic copy keyword"

This reverts commit ff53e25.

* change all to apply!, explicitly copy state for update

* fixup

* change copy scheme

* more rules

* fixup

* delete the manifest

* broken

Author: mcabbott

Manifest.toml

@@ -1,19 +1,17 @@
 name = "Optimisers"
 uuid = "3bd65402-5787-11e9-1adc-39752487f4e2"
 authors = ["Mike J Innes <[email protected]>"]
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-Functors = "0.1, 0.2"
-Requires = "0.5, 1"
-julia = "1.6" 
+Functors = "0.2.7"
+julia = "1.6"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

Project.toml

@@ -15,7 +15,7 @@ Optimisers.ADAMW
 Optimisers.AdaBelief
 ```
 
-```
+```@docs
 Optimisers.ClipGrad
 Optimisers.ClipNorm
 Optimisers.WeightDecay

docs/src/api.md

@@ -8,9 +8,9 @@ struct Descent{T}
   η::T
 end
 
-# Define an `apply` rule with which to update the current params
+# Define an `apply!` rule with which to update the current params
 # using the gradients
-function Optimisers.apply(o::Descent, state, m, m̄)
+function Optimisers.apply!(o::Descent, state, m, m̄)
     o.η .* m̄, state
 end
 
@@ -52,3 +52,7 @@ tree formed by the model and update the parameters using the gradients. Optimise
 work with different forms of gradients, but most likely use case are the gradients as
 returned by [Zygote.jl](https://fluxml.ai/Zygote.jl).
 
+There is also `Optimisers.update!` which similarly returns a new model and new state,
+but is free to mutate arrays within the old one, for efficiency.
+The method of `apply!` you write is likewise free to mutate arrays within its state;
+they are defensively copied when this rule is used with `update`.

docs/src/index.md

@@ -1,4 +1,3 @@
-patch(x, x̄) = x .- x̄
 
 function state(o, x)
   if isnumeric(x)
@@ -11,27 +10,68 @@ function state(o, x)
   end
 end
 
-function _update(o, st, x, x̄s...)
-  st′, x̄′ = apply(o, st, x, x̄s...)
-  return st′, patch(x, x̄′)
+patch!(x, x̄) = iswriteable(x) ? (x .= x .- x̄) : (x .- x̄)
+
+function _update!(o, st, x, x̄s...)
+  st′, x̄′ = apply!(o, st, x, x̄s...)
+  return st′, patch!(x, x̄′)
 end
 
-function update(o, state, x::T, x̄s...) where T
+function update!(o, state, x, x̄s...)
   if all(isnothing, x̄s)
     return state, x
   elseif isnumeric(x)
-    return _update(o, state, x, x̄s...)
+    return _update!(o, state, x, x̄s...)
   else
     x̄s′ = map(x̄ -> functor(typeof(x), x̄)[1], x̄s)
     x′, re = functor(typeof(x), x)
-    xstate = map((stᵢ, xᵢ, x̄sᵢ...) -> update(o, stᵢ, xᵢ, x̄sᵢ...), state, x′, x̄s′...)
+    xstate = map((stᵢ, xᵢ, x̄sᵢ...) -> update!(o, stᵢ, xᵢ, x̄sᵢ...), state, x′, x̄s′...)
     return map(first, xstate), re(map(last, xstate))
   end
 end
 
+function update(o, state, x, x̄s...)
+  state′ = fmap(copy, state; exclude = iswriteable)
+  x′ = fmap(copy, x; exclude = iswriteable)
+  update!(o, state′, x′, x̄s...)
+end
+
 # default all rules to first order calls
-apply(o, state, x, dx, dxs...) = apply(o, state, x, dx)
+apply!(o, state, x, dx, dxs...) = apply!(o, state, x, dx)
 
 isnumeric(x::AbstractArray{<:Number}) = isleaf(x)  # isleaf to allow for e.g. transposed shared weights
 isnumeric(x::AbstractArray{<:Bool}) = false  # convention of ChainRules is that Bool is non-differentiable
 isnumeric(x) = false
+
+iswriteable(::DenseArray{<:AbstractFloat}) = true  # more elaborate versions are possible, wait until needed?
+iswriteable(_) = false
+
+"""
+    @.. x = x + y
+    @.. x + y / z
+
+Magic broadcasting macro, for use in `apply!` rules:
+* Applied to assignment `x = ...` it is like `@.` unless `!iswriteable(x)`,
+  in which case it ignores `x`, and applies `@.` on the right.
+* Applied to other expressions, it broadcasts like `@.` but does not materialise,
+  returning a `Broadcasted` object for later use.
+"""
+macro var".."(ex)
+  if Meta.isexpr(ex, :(=))
+    dst = esc(ex.args[1])
+    src = esc(Broadcast.__dot__(ex.args[2]))
+    :(if $iswriteable($dst)
+        $dst .= $src
+      else
+        $src
+      end)
+  else
+    bc = esc(Broadcast.__dot__(ex))
+    :($lazy.($bc))
+  end
+end
+
+function lazy end
+Broadcast.broadcasted(::typeof(lazy), x) = Lazy(x)
+struct Lazy{T}; bc::T; end
+Broadcast.materialize(x::Lazy) = Broadcast.instantiate(x.bc)