Higher order forward rule for broadcast

Author: Keno

src/stage1/broadcast.jl

@@ -1,5 +1,33 @@
 using Base.Broadcast
-using Base.Broadcast: broadcasted
+using Base.Broadcast: broadcasted, Broadcasted
+
+# Forward mode broadcast rule
+struct FwdBroadcast{N, T<:AbstractTangentBundle{N}}
+  f::T
+end
+(f::FwdBroadcast{N})(args::AbstractTangentBundle{N}...) where {N} = ∂☆{N}()(f.f, args...)
+
+n_getfield(∂ₙ::∂☆{N}, b::ATB{N}, x::Union{Symbol, Int}) where {N} = ∂ₙ(ZeroBundle{N}(getfield), b, ZeroBundle{N}(x))
+
+function (∂ₙ::∂☆{N})(zc::ZeroBundle{N, typeof(copy)},
+                     bc::ATB{N, <:Broadcasted}) where {N}
+  bc = ∂ₙ(ZeroBundle{N}(Broadcast.flatten), bc)
+  args = n_getfield(∂ₙ, bc, :args)
+  r = copy(Broadcasted(
+      FwdMap(n_getfield(∂ₙ, bc, :f)),
+      ntuple(length(primal(args))) do i
+        val = n_getfield(∂ₙ, args, i)
+        if ndims(primal(val)) == 0
+          return Ref(∂ₙ(ZeroBundle{N}(getindex), val))
+        else
+          return unbundle(val)
+        end
+      end))
+  if isa(r, AbstractArray)
+    r = rebundle(r)
+  end
+  return r
+end
 
 # Broadcast over one element is just map
 function (∂⃖ₙ::∂⃖{N})(::typeof(broadcasted), f, a::Array) where {N}

src/tangent.jl

@@ -223,14 +223,23 @@ end
         $(Expr(:splatnew, B, :x))
     end
 
+expand_singleton_to_array(asize, a::AbstractZero) = fill(a, asize...)
+expand_singleton_to_array(asize, a::AbstractArray) = a
+
 function unbundle(atb::TangentBundle{Order, A}) where {Order, Dim, T, A<:AbstractArray{T, Dim}}
-    StructArray{TangentBundle{Order, T}}((atb.primal, atb.partials...))
+    asize = size(atb.primal)
+    StructArray{TangentBundle{Order, T}}((atb.primal, map(a->expand_singleton_to_array(asize, a), atb.partials)...))
 end
 
 function StructArrays.staticschema(::Type{<:TangentBundle{N, B, T}}) where {N, B, T}
     Tuple{B, T.parameters...}
 end
 
+function StructArrays.component(m::TangentBundle{N, B, T}, i::Int) where {N, B, T}
+    i == 1 && return m.primal
+    return m.partials[i - 1]
+end
+
 function StructArrays.createinstance(T::Type{<:TangentBundle}, args...)
     T(first(args), Base.tail(args))
 end
@@ -251,6 +260,10 @@ function StructArrays.staticschema(::Type{<:TaylorBundle{N, B}}) where {N, B}
     Tuple{B, Vararg{Any, N}}
 end
 
+function StructArrays.component(m::TaylorBundle{N, B}, i::Int) where {N, B, T}
+    i == 1 && return m.primal
+    return m.coeffs[i - 1]
+end
 
 function StructArrays.createinstance(T::Type{<:TaylorBundle}, args...)
     T(first(args), Base.tail(args))

test/runtests.jl

@@ -114,4 +114,13 @@ end
 # Regression tests
 @test gradient(x -> sum(abs2, x .+ 1.0), zeros(3))[1] == [2.0, 2.0, 2.0]
 
+const fwd = Diffractor.PrimeDerivativeFwd
+const bwd = Diffractor.PrimeDerivativeFwd
+
+function f_broadcast(a)
+    l = a / 2.0 * [[0. 1. 1.]; [1. 0. 1.]; [1. 1. 0.]]
+    return sum(l)
+end
+@test fwd(f_broadcast)(1.0) == bwd(f_broadcast)(1.0)
+
 include("pinn.jl")