EnzymeAD / EnzymeAD/Enzyme.jl

Structured activity annotations.

Open
#100 0 comments 1 reaction 0 assignees View on GitHub
ABI enhancement upstream
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

It could be very handy for users to support annotations nested in `NamesTuple`s/structs. Let's say we have

```julia
foo(a::Real, b::NamedTuple) = exp(log(a) + log(sum(b.x)) + log(b.y.q * b.y.r))

foo(4.2, (x = [4.0, 5.0], y = (q = 0.1, r = 0.2)))
```

and let's assume the user knows some parts of this are const. This currently doesn't work:

```julia
using Enzyme
a = Active(4.2)
b = (x = Const([4.0, 5.0]), y = (q = Active(0.1), r = Active(0.2)))
autodiff(foo, a, b)
```

```
ERROR: return type is Union{}, giving up.
```

But Enzyme is capable of handling nested mixed const/active/duplicated structures internally, of course. So the semantic equivalent

```julia
expanded_foo(a, b_x, b_y_q, b_y_r) = foo(a, (x = b_x, y = (q = b_y_q, r = b_y_r)))
da, d_y_q, d_y_r = autodiff(expanded_foo, a, b.x, b.y.q, b.y.r)
(a = da, y = (q = d_y_q, r = d_y_r))
```

works fine - it's just a bit hard on the user. :-)

Could `autodiff(foo, a, b)` be supported by just making some limited changes to Enzyme's "input formatting" stage?

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.