Implement `Reactant.batch` function for better batching (vmap too!)
#180 aperta il 16 ott 2024
Metriche repository
- Star
- (336 star)
- Metriche merge PR
- (Merge medio 2g 10h) (36 PR mergiate in 30 g)
Descrizione
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:
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.
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:
f = @compile Reactant.batch(transpose, X; dims=3)