patrick-kidger / patrick-kidger/diffrax
Better handling of errors on GPU
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.1k
- Forks
- 189
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 1
Description
When jit-compiling on the GPU: errors are currently only printed to stdout/stderr rather than properly raising. See https://github.com/google/jax/issues/9457, including the bare-bones of a possible workaround. In practice it'd be better to format the traceback appropriately, using the traceback module; see also https://github.com/google/jax/blob/7a6986c4c8fd8469bae36306efb0417b0a2f6d8c/jax/_src/traceback_util.py for inspiration.
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 with the linked JAX issue 9457 and the referenced jax/_src/traceback_util.py file to understand the proposed workaround and traceback formatting. Trace how GPU JIT errors are currently surfaced in diffrax, then verify that failures raise properly and produce an appropriate traceback instead of only printing to stdout or stderr.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100