|  | 
| 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