JuliaGPU / JuliaGPU/KernelAbstractions.jl
User-facing API like vmap
- Dominant language
- Julia
- Stars
- 523
- Forks
- 88
- Avg merge
- 1d 11h
- Merged PRs (30d)
- 25
Description
@ChrisRackauckas suggests that this package provides much of the utilities that would make broadcasting over specified axes efficient. This can be seen in [DiffEqGPU.jl](https://github.com/SciML/DiffEqGPU.jl/blob/master/src/DiffEqGPU.jl).
Can we discuss a user facing API so we can directly compare against [JAX vmap](https://jax.readthedocs.io/en/latest/jax.html?highlight=vmap#jax.vmap).
For instance if I have a function
```
f(x::Scalar, y::Vector, A::Array) = linalg...
```
How can I efficiently broadcast over collections of inputs stored in collections with axes like multidimensional arrays ("tensors").
```
# Broadcast over rows of second argument
vmap(f, in_axes=(nothing, 1, nothing))(scalar, array, array)
# Broadcast over axes for all arguments
vmap(f, in_axes=(1, 1, 3))(vector, array, tensor)
```
Further, is it possible to provide these as defaults for something like `eachslice` so that broadcasting Just Works?
```
f.(scalar, eachrow(array), array)
```
Contributor guide
No contributing guide indexed for this repository
Research direction
No target file or test is named. Start by reviewing the linked DiffEqGPU.jl implementation and the JAX vmap documentation, then inspect how KernelAbstractions.jl currently handles axes and eachslice. Done requires an agreed user-facing API, defined defaults, and tests demonstrating the proposed scalar, row, and multidimensional cases.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- hpc
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100