Skip to content
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.

Sign up for GitHub

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

merge cumprod rules to fix ambig on julia 1.6 #784

Closed
wants to merge 3 commits into from
Closed

merge cumprod rules to fix ambig on julia 1.6 #784

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
16a604d
merge cumprod rules
oxinabox Feb 9, 2024
a1a70cc
Make default dims=1
oxinabox Feb 9, 2024
b1d7b46
make sure we still error if the primal code errored
oxinabox Feb 9, 2024
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
45 changes: 15 additions & 30 deletions src/rulesets/Base/mapreduce.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -325,52 +325,37 @@ end
##### `cumprod`
#####

function rrule(::typeof(cumprod), x::AbstractVector{<:Real}; dims::Integer=1)
y = cumprod(x; dims=dims) # does nothing unless dims == 1
project_x = ProjectTo(x)
function cumprod_pullback_1(dy_raw)
dy = unthunk(dy_raw)
dx_thunk = InplaceableThunk(
dx -> if dims == 1
∇cumprod!(dx, x, dy, y)
else
dx .+= dy
end
,
@thunk project_x(if dims == 1
∇cumprod(x, dy, y)
else
dy
end)
)
return (NoTangent(), dx_thunk)
end
return y, cumprod_pullback_1
end

function rrule(::typeof(cumprod), x::AbstractArray{<:Real}; dims::Integer)
function rrule(::typeof(cumprod), x::AbstractArray{<:Real}; dims::Integer=1)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it OK to put a default value for dims even though there's no default value in base? https://github.com/JuliaLang/julia/blob/63e95d47e692352fac5541b6ee835b573f4c4898/base/accumulate.jl#L193

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, it is see:
#784 (comment)

oxinabox marked this conversation as resolved.
Show resolved Hide resolved
y = cumprod(x; dims=dims)
project_x = ProjectTo(x)
function cumprod_pullback_2(dy_raw)
function cumprod_pullback(dy_raw)
dy = unthunk(dy_raw)
dx_thunk = InplaceableThunk(
dx -> if dims <= ndims(x)
vald = Val(Int(dims))
∇cumprod_dim!(dx, vald, x, dy, y)
if ndims(x) == 1
∇cumprod!(dx, x, dy, y)
else
vald = Val(Int(dims))
∇cumprod_dim!(dx, vald, x, dy, y)
end
else
dx .+= dy
end
,
@thunk project_x(if dims <= ndims(x)
vald = Val(Int(dims))
∇cumprod_dim(vald, x, dy, y)
if ndims(x) == 1
∇cumprod(x, dy, y)
else
vald = Val(Int(dims))
∇cumprod_dim(vald, x, dy, y)
end
else
dy
end)
Comment on lines 349 to 360
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
end
,
@thunk project_x(if dims <= ndims(x)
vald = Val(Int(dims))
∇cumprod_dim(vald, x, dy, y)
if ndims(x) == 1
∇cumprod(x, dy, y)
else
vald = Val(Int(dims))
∇cumprod_dim(vald, x, dy, y)
end
else
dy
end)
end,
@thunk project_x(
if dims <= ndims(x)
if ndims(x) == 1
∇cumprod(x, dy, y)
else
vald = Val(Int(dims))
∇cumprod_dim(vald, x, dy, y)
end
else
dy
end,

)
return (NoTangent(), dx_thunk)
end
return y, cumprod_pullback_2
return y, cumprod_pullback
end

function ∇cumprod_dim(vald::Val{dim}, x::AbstractArray, dy=fill!(zero(x),1), y=cumprod(x; dims=dim)) where {dim}
Expand Down
Loading