Handle Structs with undef fields

Author: oxinabox

src/tangent_types/abstract_zero.jl

@@ -113,7 +113,11 @@ zero_tangent(x::Number) = zero(x)
 @generated function zero_tangent(primal)
     has_mutable_tangent(primal) || return ZeroTangent()  # note this takes care of tuples
     zfield_exprs = map(fieldnames(primal)) do fname
-        fval = Expr(:call, zero_tangent, Expr(:call, getfield, :primal, QuoteNode(fname)))
+        fval = if isdefined(primal, fname)
+            Expr(:call, zero_tangent, Expr(:call, getfield, :primal, QuoteNode(fname)))
+        else
+            ZeroTangent()
+        end
         Expr(:kw, fname, fval)
     end
     backing_expr = Expr(:tuple, Expr(:parameters, zfield_exprs...))

test/tangent_types/abstract_zero.jl

@@ -180,7 +180,7 @@ end
     @test zero_tangent([1.0, 2.0]) == [0.0, 0.0]
     @test zero_tangent([[1.0, 2.0], [3.0]]) == [[0.0, 0.0], [0.0]]
 
-    @testset "undef elements" begin
+    @testset "undef elements Vector" begin
         x = Vector{Vector{Float64}}(undef, 3)
         x[2] = [1.0, 2.0]
         dx = zero_tangent(x)
@@ -202,4 +202,20 @@ end
         @test length(db) == 3
         @test db isa Vector
     end
+
+    @testset "undef fields struct" begin
+        dx = zero_tangent(Core.Box())
+        @test dx.contents isa ZeroTangent
+        @test (dx.contents = 2.0) == 2.0  # should be assignable
+
+        mutable struct MyPartiallyDefinedStruct
+            intro::Float64
+            contents::Number
+            MyPartiallyDefinedStruct(x) = new(x)
+        end
+        dy = zero_tangent(MyPartiallyDefinedStruct(1.5))
+        @test iszero(dy.intro)
+        @test iszero(dy.contents)
+        @test (dy.contents = 2.0) == 2.0  # should be assignable
+    end
 end