Lightning-AI / Lightning-AI/LitServe
auto accelerator for JAX
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.9k
- Forks
- 304
- Avg merge
- 3d 13h
- Merged PRs (30d)
- 6
Description
Creating this issue to refer for implementing auto accelerator for JAX users.
The thing is that we need to be opinionated about how other frameworks decide to support devices and what devices they support. So the semantics of `auto` will need to be framework-specific.
For instance with JAX you can call device_put on GPU, but you don't have mps. So we can add JAX with a different implementation for _choose_gpu_accelerator_backend.
We should actually have framework-specific _choose_gpu_accelerator_backend functions, like _choose_gpu_accelerator_torch etc.
Originally posted by @lantiga in https://github.com/Lightning-AI/litserve/pull/44#discussion_r1572349925
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by locating the existing _choose_gpu_accelerator_backend implementation and the framework-specific accelerator selection logic. Read how current frameworks handle auto device selection, then define the JAX-specific path and its supported devices without assuming MPS. Done means JAX users can use auto accelerator selection with framework-appropriate semantics.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100