google-research / google-research/dex-lang

Support AD on ADTs that have tangent types equal to the primal types

Open
#565 2 comments 1 reaction 0 assignees View on GitHub
language / autodiff language / compiler
Dominant language
Haskell
Stars
1.7k
Forks
116
PR merge metrics
No merged PRs in 30d

Description

Currently, `grad` does not work on data types such as `List`. For example:
```
ls = (AsList _ [1.,2.,3.])
def sumList ((AsList _ arr) : List Float): Float = sum arr
sumList ls
> 6.
(grad sumList) ls
dex: Not implemented
CallStack (from HasCallStack):
error, called at src/lib/Autodiff.hs:507:18 in dex-0.1.0.0-RKG8rO926ZDrlxlc3NaeB:Autodiff
notImplemented, called at src/lib/Autodiff.hs:386:19 in dex-0.1.0.0-RKG8rO926ZDrlxlc3NaeB:Autodiff
```
Discussion on Friday @apaszke @danieldjohnson @duvenaud
For some data types, taking `grad` with respect to them do not require user-supplied tangent types, since they are the same as the original type. This would be easier to implement than to support `grad` on any, general data structure.

Contributor guide

Open the contributing guide

Research direction

Start with the failing example using List, AsList, and grad, then inspect src/lib/Autodiff.hs at the reported locations around lines 386 and 507. Trace how tangent types are handled for data types whose tangent type equals the primal type. Done means the example differentiates without the current Not implemented error.

Written by the indexing model from the issue text.

Assessment

Tech stack
haskell
Domain
compilers
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 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.