Enzyme.jl/test/ext/staticarrays.jl at e6b5c3acc6cf7b85319859abce030710afd77d2c · EnzymeAD/Enzyme.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
using Enzyme, Test

using StaticArrays

@testset "Gradient & StaticArrays" begin

    x = @SArray [5.0 0.0 6.0]
    dx = Enzyme.gradient(Reverse, prod, x)[1]
    @test dx isa SArray
    @test dx ≈ [0 30 0]

    x = @SVector [1.0, 2.0, 3.0]
    y = onehot(x)
    # this should be a very specific type of SArray, but there
    # is a bizarre issue with older julia versions where it can be MArray
    @test eltype(y) <: StaticVector
    @test length(y) == 3
    @test y[1] == [1.0, 0.0, 0.0]
    @test y[2] == [0.0, 1.0, 0.0]
    @test y[3] == [0.0, 0.0, 1.0]

    y = onehot(x, 2, 3)
    @test eltype(y) <: StaticVector
    @test length(y) == 2
    @test y[1] == [0.0, 1.0, 0.0]
    @test y[2] == [0.0, 0.0, 1.0]

    x = @SArray [5.0 0.0 6.0]
    dx = Enzyme.gradient(Forward, prod, x)[1]
    @test dx isa SArray
    @test dx ≈ [0 30 0]

    f0 = x -> sum(2*x)
    f1 = x -> @SVector Float64[x[2], 2*x[2]]
    f2 = x -> @SMatrix Float64[x[2] x[1]; 2*x[2] 2*x[1]]

    x = @SVector Float64[1, 2]

    @inferred gradient(Forward, f0, x)
    dx = gradient(Forward, f0, x)[1]
    @test dx isa SVector
    @test dx == [2.0, 2.0]  # test to make sure conversion works
    @test gradient(Forward, f1, x)[1] isa SMatrix
    @test gradient(Forward, f1, x)[1] == [0 1.0; 0 2.0]
    @test Enzyme.jacobian(Forward, f2, x)[1] isa SArray
    @test Enzyme.jacobian(Forward, f2, x)[1] == reshape(Float64[0,0,1,2,1,2,0,0], (2,2,2))

    x = @SMatrix Float64[1 2; 3 4]

    @inferred gradient(Forward, f0, x)
    dx = gradient(Forward, f0, x)[1]
    @test dx isa SMatrix
    @test dx == fill(2.0, (2,2))
    @test gradient(Forward, f1, x)[1] isa SArray
    @test gradient(Forward, f1, x)[1] == reshape(Float64[0,0,1,2,0,0,0,0], (2,2,2))
    @test Enzyme.jacobian(Forward, f2, x)[1] isa SArray
    @test Enzyme.jacobian(Forward, f2, x)[1] == reshape(
        Float64[0,0,1,2,1,2,0,0,0,0,0,0,0,0,0,0], (2,2,2,2),
    )

    x = @SVector Float64[1, 2]

    @inferred gradient(Reverse, f0, x)
    dx = gradient(Reverse, f0, x)[1]
    @test dx isa SVector
    @test dx == [2.0, 2.0]  # test to make sure conversion works
    @test_broken gradient(Reverse, f1, x)[1] isa SMatrix
    @test_broken gradient(Reverse, f1, x)[1] == [0 1.0; 0 2.0]
    @test_broken Enzyme.jacobian(Reverse, f2, x)[1] isa SArray
    @test_broken Enzyme.jacobian(Reverse, f2, x)[1] == reshape(Float64[0,0,1,2,1,2,0,0], (2,2,2))

    x = @SMatrix Float64[1 2; 3 4]

    @test_broken gradient(Reverse, f1, x)[1] isa SArray
    @test_broken gradient(Reverse, f1, x)[1] == reshape(Float64[0,0,1,2,0,0,0,0], (2,2,2))
    @test_broken Enzyme.jacobian(Reverse, f2, x)[1] isa SArray
    @test_broken Enzyme.jacobian(Reverse, f2, x)[1] == reshape(
        Float64[0,0,1,2,1,2,0,0,0,0,0,0,0,0,0,0], (2,2,2,2),
    )
end

function unstable_fun(A0)
    A = 'N' in ('H', 'h', 'S', 's') ? wrap(A0) : A0
    (@inbounds A[1])::eltype(A0)
end
@testset "Type unstable static array index" begin
    inp = ones(SVector{2, Float64})
    res = Enzyme.gradient(Enzyme.Reverse, unstable_fun, inp)[1]
    @test res ≈ [1.0, 0.0]
    res = Enzyme.gradient(Enzyme.Forward, unstable_fun, inp)[1]
    @test res ≈ [1.0, 0.0]
end

function inner_forhess(x)
    return tanh.(x)
end

function for_hess(x)
    return sum(inner_forhess(x))
end

grad_forhess(x) = autodiff(Reverse, for_hess, Active, Active(x))[1][1]
hess(x) = jacobian(Forward, grad_forhess, x)[1]

@testset "StaticArrays hessian" begin
    x = @SVector zeros(10)
    res = [-2.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 -2.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 -2.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 -2.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 -2.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 -2.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 -2.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 -2.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 -2.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 -2.0]
    @test jacobian(Forward, grad_forhess, x)[1] ≈ res
end