EnzymeAD / EnzymeAD/Reactant.jl

Implement `Reactant.batch` function for better batching (vmap too!)

Open
#180 0 comments 0 reactions 0 assignees View on GitHub
enhancement good first issue
Dominant language
Julia
Stars
370
Forks
74
Avg merge
18h 47m
Merged PRs (30d)
30

Description

Currently, the way batching is implemented is by replacing broadcasting with `enzyme.batch` op and tracing over the broadcasted code. The following example should just work:

```julia
X = [rand(4,4) for _ in 1:10]
f = @compile transpose.(X)
```

One inconvenient of Julia's broadcasting is that there is no way to specify the dimension over which to broadcast; it will just iterate over everything. Thus, users need to use `eachslice` for slicing over the desired dimension.

```julia
X = rand(4,4,10)
f = @compile broadcast(transpose, eachslice(X, dims=3))
```

I'm not sure if we would correctly then batch on the desired dimension in this case or that it would create some extra instructions... need to check it.

But it could be beneficial to have some similar functionality in one `batch` function which would be easier to correctly trace and users coming from Jax would be more familiarized. An example:

```julia
f = @compile Reactant.batch(transpose, X; dims=3)
```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.