JuliaGPU / JuliaGPU/KernelAbstractions.jl

User-facing API like vmap

Open
#117 11 comments 2 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.