diff --git a/src/tracked.jl b/src/tracked.jl index 50c77c4e..827e00ea 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -473,6 +473,10 @@ end Base.one(::Type{TrackedReal{V,D,O}}) where {V,D,O} = TrackedReal{V,D,O}(one(V)) Base.zero(::Type{TrackedReal{V,D,O}}) where {V,D,O} = TrackedReal{V,D,O}(zero(V)) +# neither identity depends on the origin (#172) +Base.one(::Type{TrackedReal{V,D}}) where {V,D} = TrackedReal{V,D,Nothing}(one(V)) +Base.zero(::Type{TrackedReal{V,D}}) where {V,D} = TrackedReal{V,D,Nothing}(zero(V)) + Base.rand(::Type{TrackedReal{V,D,O}}) where {V,D,O} = TrackedReal{V,D,O}(rand(V)) Base.rand(rng::Random.AbstractRNG, ::Type{TrackedReal{V,D,O}}) where {V,D,O} = TrackedReal{V,D,O}(rand(rng, V)) diff --git a/test/TrackedTests.jl b/test/TrackedTests.jl index 0ef5648a..bdbf8442 100644 --- a/test/TrackedTests.jl +++ b/test/TrackedTests.jl @@ -743,6 +743,14 @@ for inds in ((rowmask, :), (:, colmask), (rowmask, 2:3), (mask,), (vec(mask),)) empty!(tp) end +# a 0-d array has no index to dispatch on (#126) +tp0 = InstructionTape() +ta0 = TrackedArray(fill(varr[1]), fill(darr[1]), tp0) +@test samefields(ta0[], ta0[1]) +@test samefields(ta0[], ta0[CartesianIndex()]) +@test ta0[].origin === ta0 +@test isempty(tp0) + # `view` aliases the parent's buffers, so nothing needs to be recorded ta_view = @inferred view(ta, :, 2) @test samefields(ta_view, TrackedArray(varr[:, 2], darr[:, 2], tp)) @@ -808,6 +816,10 @@ tr_float32 = TrackedReal(Float32(v_float), Float32(d), tp) @test samefields(zero(tr_float), typeof(tr_float)(zero(v_float))) +# an eltype whose origin is free still has both identities (#172) +@test samefields(one(TrackedReal{Float64,Float64}), TrackedReal{Float64,Float64,Nothing}(1.0)) +@test samefields(zero(TrackedReal{Float64,Float64}), TrackedReal{Float64,Float64,Nothing}(0.0)) + tr_rand = rand(TrackedReal{Int,Float64,Nothing}) @test samefields(tr_rand, TrackedReal{Int,Float64,Nothing}(ReverseDiff.value(tr_rand))) diff --git a/test/api/GradientTests.jl b/test/api/GradientTests.jl index 3f755ca5..ec72257f 100644 --- a/test/api/GradientTests.jl +++ b/test/api/GradientTests.jl @@ -406,4 +406,49 @@ end @test ReverseDiff.gradient(x -> sum([Real[x[1], x[2]] x[3]*ones(2)]), [1.0, 2.0, 3.0]) == [1.0, 1.0, 2.0] end +############################################################################################ + +@testset "indexing a zero-dimensional `TrackedArray` (#126)" begin + # the MWE from the issue + @test ReverseDiff.gradient(x -> x[], fill(1.0)) == fill(1.0) + + # `t[]`, `t[1]` and `t[CartesianIndex()]` are the same element + @test ReverseDiff.gradient(x -> x[]^3, fill(2.0)) == fill(12.0) + @test ReverseDiff.gradient(x -> x[1]^3, fill(2.0)) == fill(12.0) + @test ReverseDiff.gradient(x -> x[CartesianIndex()]^3, fill(2.0)) == fill(12.0) + + # the other APIs, and tape replay + @test ReverseDiff.jacobian(x -> fill(x[]^2), fill(3.0)) == fill(6.0, 1, 1) + @test ReverseDiff.hessian(x -> x[]^3, fill(2.0)) == fill(12.0, 1, 1) + + tape = ReverseDiff.GradientTape(x -> x[]^3, fill(2.0)) + @test ReverseDiff.gradient!(tape, fill(4.0)) == fill(48.0) + @test ReverseDiff.gradient!(ReverseDiff.compile(tape), fill(4.0)) == fill(48.0) +end + +############################################################################################ + +@testset "reductions with `dims` over mixed tracked reals (#172)" begin + # `x[i]` has an origin and `2 * x[i]` has none, so the eltype leaves it free + mixed(x) = map(i -> isodd(i) ? x[i] : 2 * x[i], eachindex(x)) + + @test ReverseDiff.gradient(x -> sum(sum(mixed(x); dims = 1)), [1.0, 2.0]) == [1.0, 2.0] + @test ReverseDiff.gradient(x -> sum(prod(mixed(x); dims = 1)), [3.0, 5.0]) == [10.0, 6.0] + @test ReverseDiff.gradient(x -> sum(sum(reshape(mixed(x), 2, 2); dims = 2)), + [1.0, 2.0, 3.0, 4.0]) == [1.0, 2.0, 1.0, 2.0] + @test ReverseDiff.hessian(x -> sum(sum(mixed(x); dims = 1))^2, [1.0, 2.0]) == + [2.0 4.0; 4.0 8.0] + + # `prod` is nonlinear, so a replay reusing stale values would be caught here + tape = ReverseDiff.GradientTape(x -> sum(prod(mixed(x); dims = 1)), [3.0, 5.0]) + @test ReverseDiff.gradient!(tape, [5.0, 6.0]) == [12.0, 10.0] + @test ReverseDiff.gradient!(ReverseDiff.compile(tape), [5.0, 6.0]) == [12.0, 10.0] + + # an abstract eltype reaches a different reduction path + @test ReverseDiff.gradient(x -> sum(sum(Real[x[1], 2 * x[2]]; dims = 1)), [1.0, 2.0]) == + [1.0, 2.0] + @test ReverseDiff.gradient(x -> sum(sum(Any[x[1], 2 * x[2]]; dims = 1)), [1.0, 2.0]) == + [1.0, 2.0] +end + end # module