Skip to content

Use RealDot.realdot #542

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Oct 24, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
name = "ChainRules"
uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
version = "1.11.6"
version = "1.12.0"

[deps]
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
RealDot = "c1ae055f-0cd5-4b69-90a6-9a35b1a98df9"
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[compat]
Expand All @@ -15,6 +16,7 @@ ChainRulesTestUtils = "1"
Compat = "3.35"
FiniteDifferences = "0.12.8"
JuliaInterpreter = "0.8"
RealDot = "0.1"
StaticArrays = "1.2"
julia = "1"

Expand Down
1 change: 1 addition & 0 deletions src/ChainRules.jl
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ using Compat
using LinearAlgebra
using LinearAlgebra.BLAS
using Random
using RealDot: realdot
using Statistics

# Basically everything this package does is overloading these, so we make an exception
Expand Down
2 changes: 1 addition & 1 deletion src/rulesets/Base/base.jl
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ end

function frule((_, Δz), ::typeof(hypot), z::Complex)
Ω = hypot(z)
∂Ω = _realconjtimes(z, Δz) / ifelse(iszero(Ω), one(Ω), Ω)
∂Ω = realdot(z, Δz) / ifelse(iszero(Ω), one(Ω), Ω)
return Ω, ∂Ω
end

Expand Down
6 changes: 3 additions & 3 deletions src/rulesets/Base/fastmath_able.jl
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ let
Ω = abs(x)
# `ifelse` is applied only to denominator to ensure type-stability.
signx = x isa Real ? sign(x) : x / ifelse(iszero(x), one(Ω), Ω)
return Ω, _realconjtimes(signx, Δx)
return Ω, realdot(signx, Δx)
end

function rrule(::typeof(abs), x::Union{Real, Complex})
Expand All @@ -82,7 +82,7 @@ let

## abs2
function frule((_, Δz), ::typeof(abs2), z::Union{Real, Complex})
return abs2(z), 2 * _realconjtimes(z, Δz)
return abs2(z), 2 * realdot(z, Δz)
end

function rrule(::typeof(abs2), z::Union{Real, Complex})
Expand Down Expand Up @@ -146,7 +146,7 @@ let
) where {T<:Union{Real,Complex}}
Ω = hypot(x, y)
n = ifelse(iszero(Ω), one(Ω), Ω)
∂Ω = (_realconjtimes(x, Δx) + _realconjtimes(y, Δy)) / n
∂Ω = (realdot(x, Δx) + realdot(y, Δy)) / n
return Ω, ∂Ω
end

Expand Down
6 changes: 3 additions & 3 deletions src/rulesets/Base/mapreduce.jl
Original file line number Diff line number Diff line change
Expand Up @@ -111,13 +111,13 @@ function frule(
ẋ = unthunk(Δx)
y = sum(abs2, x; dims=dims)
∂y = if dims isa Colon
2 * real(dot(x, ẋ))
2 * realdot(x, ẋ)
elseif VERSION ≥ v"1.2" # multi-iterator mapreduce introduced in v1.2
mapreduce(+, x, ẋ; dims=dims) do xi, dxi
2 * _realconjtimes(xi, dxi)
2 * realdot(xi, dxi)
end
else
2 * sum(_realconjtimes.(x, ẋ); dims=dims)
2 * sum(realdot.(x, ẋ); dims=dims)
end
return y, ∂y
end
Expand Down
7 changes: 0 additions & 7 deletions src/rulesets/Base/utils.jl
Original file line number Diff line number Diff line change
@@ -1,10 +1,3 @@
# real(conj(x) * y) avoiding computing the imaginary part if possible
@inline _realconjtimes(x, y) = real(conj(x) * y)
@inline _realconjtimes(x::Complex, y::Complex) = muladd(real(x), real(y), imag(x) * imag(y))
@inline _realconjtimes(x::Real, y::Complex) = x * real(y)
@inline _realconjtimes(x::Complex, y::Real) = real(x) * y
@inline _realconjtimes(x::Real, y::Real) = x * y

# imag(conj(x) * y) avoiding computing the real part if possible
@inline _imagconjtimes(x, y) = imag(conj(x) * y)
@inline function _imagconjtimes(x::Complex, y::Complex)
Expand Down
4 changes: 2 additions & 2 deletions src/rulesets/LinearAlgebra/blas.jl
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ function frule((_, Δx), ::typeof(BLAS.nrm2), x)
∂Ω = if x isa Real
BLAS.dot(x, Δx) / s
else
sum(y -> _realconjtimes(y...), zip(x, Δx)) / s
sum(y -> realdot(y...), zip(x, Δx)) / s
end
return Ω, ∂Ω
end
Expand Down Expand Up @@ -72,7 +72,7 @@ end

function frule((_, Δx), ::typeof(BLAS.asum), x)
∂Ω = sum(zip(x, Δx)) do (xi, Δxi)
return _realconjtimes(_signcomp(xi), Δxi)
return realdot(_signcomp(xi), Δxi)
end
return BLAS.asum(x), ∂Ω
end
Expand Down
2 changes: 1 addition & 1 deletion src/rulesets/LinearAlgebra/factorization.jl
Original file line number Diff line number Diff line change
Expand Up @@ -347,7 +347,7 @@ function _eigen_norm_phase_fwd!(∂V, A, V)
@inbounds for i in axes(V, 2)
v, ∂v = @views V[:, i], ∂V[:, i]
# account for unit normalization
∂c_norm = -real(dot(v, ∂v))
∂c_norm = -realdot(v, ∂v)
if eltype(V) <: Real
∂c = ∂c_norm
else
Expand Down
6 changes: 3 additions & 3 deletions src/rulesets/LinearAlgebra/norm.jl
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ function frule((_, ẋ), ::typeof(norm), x::Number, p::Real)
zero(real(x)) * zero(real(Δx))
else
signx = x isa Real ? sign(x) : x * pinv(y)
_realconjtimes(signx, Δx)
realdot(signx, Δx)
end
return y, ∂y
end
Expand Down Expand Up @@ -235,7 +235,7 @@ function rrule(::typeof(LinearAlgebra.norm2), x::AbstractArray{<:Number})
end

function _norm2_forward(x, Δx, y)
∂y = real(dot(x, Δx)) * pinv(y)
∂y = realdot(x, Δx) * pinv(y)
return ∂y
end
function _norm2_back(x, y, Δy)
Expand Down Expand Up @@ -280,7 +280,7 @@ function rrule(::typeof(normalize), x::AbstractVector{<:Number})
LinearAlgebra.__normalize!(y, nrm)
function normalize_pullback(ȳ)
Δy = unthunk(ȳ)
∂x = (Δy .- real(dot(y, Δy)) .* y) .* pinv(nrm)
∂x = (Δy .- realdot(y, Δy) .* y) .* pinv(nrm)
return (NoTangent(), ∂x)
end
normalize_pullback(::ZeroTangent) = (NoTangent(), ZeroTangent())
Expand Down
2 changes: 1 addition & 1 deletion src/rulesets/LinearAlgebra/symmetric.jl
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,7 @@ function frule(
# diag(U' * tmp) without computing matrix product
∂λ = similar(λ)
@inbounds for i in eachindex(λ)
∂λ[i] = @views real(dot(U[:, i], tmp[:, i]))
∂λ[i] = @views realdot(U[:, i], tmp[:, i])
end
return λ, ∂λ
end
Expand Down