EnzymeAD / EnzymeAD/Enzyme.jl

API for multi-arg hvp!/hvp

Open
#2,853 5 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

Right now, it’s not easy (impossible?) to use annotations like `Const` with `hvp!`.

Use case:

Suppose we have a function:

```julia
f(x, y, z) = sum(map(sin, x)) + sum(map(cos, y)) + sum(map(sin, z))
```

It would be nice to support calls like:
```Julia
hvp(
f,
[1.0, 2.0],
Const([1.0, 2.0]),
[1.0, 2.0],
[1.0, 1.0, 1.0, 1.0],
)
```
Open question: how should directions be passed?

If hvp takes a tuple of arguments, should it also take a tuple of direction vectors, or is a single “flattened” direction vector of the correct total dimension sufficient?

Interface options

Option A: one direction per active argument
• For one active argument: `hvp(f, x, Const(y), vx)`
• For two active arguments: `hvp(f, x, y, vx, vy)`

Option B: a single combined direction vector

Pass a single container holding all directions (for example, a vector/tuple of vectors): `hvp(f, x, y, [vx, vy])`

The same question applies to output buffers.

Whichever convention we pick for direction vectors should likely also apply to the output buffers (e.g., one buffer per active argument vs. one combined buffer).

I guess someone can propose alternative options.

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.