EnzymeAD / EnzymeAD/Reactant.jl
Implement `Reactant.batch` function for better batching (vmap too!)
- 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.