`sum` rules do not match the signatures of the Base methods

Open
  • Is something wrong?
#868 1 comment 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

  • #870 by @Huyvux12 — closed without merging

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
35/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Stale
Domain
backend

Research direction

Read the sum and sum! rules in src/rulesets/Base/mapreduce.jl, including the existing foldl rule's treatment of init. Update the rules to match the Base signatures and keywords described in the issue, then check that unsupported generic iterables still fall back to AD. Pull request #870 is already open, so this work is under way.

Written by the indexing model from the issue text.

Description

The sum rules don't match the signatures of the Base methods they cover. If a call matches a rule's positional types but not its keywords, the rule throws a MethodError instead of returning nothing and leaving the call to AD. This is the same kind of bug as #768.

Base methods (same on Julia 1.10 and 1.13)
sum(a; kw...)                                 # -> mapreduce -> mapfoldl: accepts `init`
sum(a::AbstractArray; dims=:, kw...)          # accepts `dims`, `init`
sum(f, a; kw...)                              # accepts `init`
sum(f, a::AbstractArray; dims=:, kw...)       # accepts `dims`, `init`
sum!(r::AbstractArray, A::AbstractArray; init=true)
sum!(f, r::AbstractArray, A::AbstractArray; init=true)

(The specialisations for ranges, Bool, BigInt and BigFloat are faster implementations of the same function. Bool is already @non_differentiable.)

Rules that are missing a keyword Base accepts

None of the sum rules accept init:

  • frule/rrule for sum(x::AbstractArray; dims=:)
  • frule/rrule for sum(x::Tuple), which covers part of Base's generic sum(a; kw...) but takes no keywords
  • frule for sum!(y, x). With init=false the output accumulates into y, so the tangent would be sum!(ẏ, ẋ; init=false)
  • rrules for sum(f, x::Tuple) and sum(f, x::AbstractArray; dims=:)
  • frule/rrule for sum(abs2, x::AbstractArray; dims=:)

Example on current main:

julia> frule((NoTangent(), ones(3)), sum, rand(3); init=0.0)
ERROR: MethodError: no method matching frule(::Tuple{NoTangent, Vector{Float64}}, ::typeof(sum), ::Vector{Float64}; init::Float64)

julia> rrule(sum, rand(3); init=0.0)
ERROR: MethodError: no method matching rrule(::typeof(sum), ::Vector{Float64}; init::Float64)
Proposal

Make every sum/sum! rule accept exactly the arguments and keywords of the Base method it covers. Treat init as non-differentiable (no tangent), as the foldl rrule already does (src/rulesets/Base/mapreduce.jl, "it does not return a gradient for init").

There can't be a rule for generic iterables: the tangent of a Generator, Dict and so on is structural rather than elementwise, so sum(ẋ) would be meaningless. Those inputs should keep falling back to AD (see #768 and #867).

Dominant language
Julia
Stars
479
Forks
100
Avg merge
1d 4h
Merged PRs (30d)
14

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from JuliaDiff/ChainRules.jl

All issues in JuliaDiff/ChainRules.jl

Similar issues

More Julia issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.