JuliaDiff / JuliaDiff/ChainRulesCore.jl

Write docs on how to write rules to run different computations at different times

Open
#46 3 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

documentation
Dominant language
Julia
Stars
267
Forks
66
PR merge metrics
No merged PRs in 30d

Description

With #30
this might be a great example

function rrule(::typeof(/), A::AbstractVecOrMat{<:Real}, B::AbstractVecOrMat{<:Real})
    # Compute C (return value differently) to just calling `A/B`
    # so we have some parts precomputed 
    # that we need to use for computing the pullback

    Aᵀ, dA_pb = rrule(adjoint, A)
    Bᵀ, dB_pb = rrule(adjoint, B)
    Cᵀ, dS_pb = rrule(\, Bᵀ, Aᵀ)
    C, dC_pb = rrule(adjoint, Cᵀ)
    
    function slash_pullback(Ȳ)
         # These two run no matter which partial you want
         # but they don't run if you only want the result (C)

        _, dC = dC_pb(Ȳ)
        _, dAᵀ, dBᵀ = dS_pb(dC)

        # Each of these only run if you want that particular partial
        ∂A = @thunk last(dA_pb(dAᵀ))
        ∂B = @thunk last(dA_pb(dBᵀ))

        (NO_FIELDS, ∂A, ∂B)
    end
    return C, slash_pullback
end

Contributor guide

No contributing guide indexed for this repository

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.

Research direction

Start by reading linked issue #30 and the rrule example in this issue, focusing on which computations run for the result versus requested partials. Done means documentation explains how to write rules with this timing behavior and makes the example understandable; no documentation file or test is named in the issue.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
documentation
Issue type
Documentation
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.