|
20 | 20 | function rrule(::typeof(norm), x::AbstractArray{<:Number}, p::Real)
|
21 | 21 | y = LinearAlgebra.norm(x, p)
|
22 | 22 | function norm_pullback_p(Δy)
|
23 |
| - ∂x = InplaceableThunk( |
| 23 | + # ∂x = InplaceableThunk( |
24 | 24 | # out-of-place versions
|
25 |
| - if isempty(x) || p == 0 |
26 |
| - @thunk(zero.(x) .* (zero(y) * zero(real(Δy)))) |
| 25 | + ∂x = @thunk(if isempty(x) || p == 0 |
| 26 | + zero.(x) .* (zero(y) * zero(real(Δy))) |
27 | 27 | elseif p == 2
|
28 |
| - @thunk(_norm2_back(x, y, Δy)) |
| 28 | + _norm2_back(x, y, Δy) |
29 | 29 | elseif p == 1
|
30 |
| - @thunk(_norm1_back(x, y, Δy)) |
| 30 | + _norm1_back(x, y, Δy) |
31 | 31 | elseif p == Inf
|
32 |
| - @thunk(_normInf_back(x, y, Δy)) |
| 32 | + _normInf_back(x, y, Δy) |
33 | 33 | elseif p == -Inf
|
34 |
| - @thunk(_normInf_back(x, y, Δy)) |
| 34 | + _normInf_back(x, y, Δy) |
35 | 35 | else
|
36 |
| - @thunk(_normp_back_x(x, p, y, Δy)) |
37 |
| - end, |
38 |
| - # in-place versions |
39 |
| - if isempty(x) || p == 0 |
40 |
| - identity |
41 |
| - elseif p == 2 |
42 |
| - dx -> _norm2_back!(dx, x, y, Δy) |
43 |
| - elseif p == 1 |
44 |
| - dx -> _norm1_back!(dx, x, y, Δy) |
45 |
| - elseif p == Inf |
46 |
| - dx -> dx .+= _normInf_back(x, y, Δy) # not really in-place! could perhaps be improved |
47 |
| - elseif p == -Inf |
48 |
| - dx -> dx .+= _normInf_back(x, y, Δy) |
49 |
| - else |
50 |
| - dx -> dx .+= _normp_back_x(x, p, y, Δy) |
51 |
| - end |
52 |
| - ) |
| 36 | + _normp_back_x(x, p, y, Δy) |
| 37 | + end) |
| 38 | + # , # in-place versions -- can be fixed when actually useful? |
| 39 | + # dx -> if isempty(x) || p == 0 |
| 40 | + # dx |
| 41 | + # elseif p == 2 |
| 42 | + # _norm2_back!(dx, x, y, Δy) |
| 43 | + # elseif p == 1 |
| 44 | + # _norm1_back!(dx, x, y, Δy) |
| 45 | + # elseif p == Inf |
| 46 | + # dx .+= _normInf_back(x, y, Δy) # not really in-place! could perhaps be improved |
| 47 | + # elseif p == -Inf |
| 48 | + # dx .+= _normInf_back(x, y, Δy) |
| 49 | + # else |
| 50 | + # dx .+= _normp_back_x(x, p, y, Δy) |
| 51 | + # end |
| 52 | + # ) |
53 | 53 | ∂p = @thunk _normp_back_p(x, p, y, Δy)
|
54 | 54 | return (NO_FIELDS, ∂x, ∂p)
|
55 | 55 | end
|
|
59 | 59 | function rrule(::typeof(norm), x::AbstractArray{<:Number})
|
60 | 60 | y = LinearAlgebra.norm(x)
|
61 | 61 | function norm_pullback_2(Δy)
|
62 |
| - ∂x = InplaceableThunk( |
63 |
| - if isempty(x) |
64 |
| - @thunk(zero.(x) .* (zero(y) * zero(real(Δy)))) |
65 |
| - else |
66 |
| - @thunk(_norm2_back(x, y, Δy)) |
67 |
| - end |
68 |
| - , |
69 |
| - if isempty(x) |
70 |
| - identity |
| 62 | + # ∂x = InplaceableThunk( |
| 63 | + ∂x = @thunk(if isempty(x) |
| 64 | + zero.(x) .* (zero(y) * zero(real(Δy))) |
71 | 65 | else
|
72 |
| - dx -> _norm2_back!(dx, x, y, Δy) |
73 |
| - end |
74 |
| - ) |
| 66 | + _norm2_back(x, y, Δy) |
| 67 | + end) |
| 68 | + # , |
| 69 | + # dx -> if isempty(x) |
| 70 | + # dx |
| 71 | + # else |
| 72 | + # _norm2_back!(dx, x, y, Δy) |
| 73 | + # end |
| 74 | + # ) |
75 | 75 | return (NO_FIELDS, ∂x)
|
76 | 76 | end
|
77 | 77 | norm_pullback_2(::Zero) = (NO_FIELDS, Zero())
|
|
0 commit comments