Skip to content

Commit

Permalink
Handle Structs with undef fields
Browse files Browse the repository at this point in the history
  • Loading branch information
oxinabox committed Dec 19, 2023
1 parent f7213c8 commit 7718bc7
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 2 deletions.
6 changes: 5 additions & 1 deletion src/tangent_types/abstract_zero.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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...))
Expand Down
18 changes: 17 additions & 1 deletion test/tangent_types/abstract_zero.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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

0 comments on commit 7718bc7

Please sign in to comment.