JuliaDiff / JuliaDiff/ReverseDiff.jl

Support undifferentiated parameters for pre-recorded API

Open
#36 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

priority
Dominant language
Julia
Stars
393
Forks
60
Avg merge
18h 24m
Merged PRs (30d)
8

Description

Right now, if I have a function f(a, b, c) and I only want to create a function which returns the gradient w.r.t. to a and b, I have two options:

  • ∇f(a, b, c) = ReverseDiff.gradient((x, y) -> f(x, y, c), (a, b)))
  • ∇f! = ReverseDiff.compile_gradient(f, (a, b, c)), and just ignore the c gradient that will pop out

The former has to re-record the function for every call, while the latter wastes some computation differentiating w.r.t. c.

We should support something akin to Tensorflow's placeholders for the pre-recorded API, allowing you to drop in updatable parameters that aren't differentiated against. This can be accomplished by recording the tape as normal, and then "turning off" differentiation on the selected parameters (the idiom for that currently is to set the tape to NULL_TAPE, but I'm going to play around with it). Some refactoring should probably be done to get the most out of this change performance-wise (e.g., allow the instantiation of a TrackedArray with deriv == nothing).

As for the API, I can think of two different paths we could take:

  • Select which arguments are to be differentiated against using a wrt function, e.g. ReverseDiff.compile_gradient(f, (wrt(a), wrt(b), c))
  • Select which arguments are not to be differentiated against using a param function, e.g. ReverseDiff.compile_gradient(f, (a, b, param(c)))

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 with ReverseDiff.compile_gradient and the existing NULL_TAPE idiom described in the issue, then inspect how TrackedArray handles deriv == nothing. Compare the proposed wrt and param API paths. Done means pre-recorded functions can update selected parameters without differentiating against them or recomputing those gradients.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
api, backend-api-design
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.