jax-ml/jax

View on GitHub

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Stars: 36.3k
Forks: 3.8k
Open beginner issues: 12
Indexed issues: 1.8k
Avg merge: 1d 6h
Merged PRs (30d): 399

Dominant language: Python
License: Apache-2.0
Last GitHub push: Sep 16, 2026
Latest indexed: Sep 16, 2026
Contributing guide: Contributing guide
Code of conduct: No code of conduct
Beginner labels: good first issue

Open All

1749 open issues indexed

Linearity (or affineness?) checking for PRNG keys Open

enhancement

jax-ml/jax#192 · 8 comments · 7 reactions · 1 assignee · Indexed Sep 16, 2026
jax.numpy array indexing has different out-of-bounds behavior to numpy Open

documentation NVIDIA GPU P2 (eventual)

jax-ml/jax#278 · 16 comments · 7 reactions · 1 assignee · Indexed Sep 16, 2026
jit cache leaks memory Open

bug

jax-ml/jax#282 · 4 comments · 1 reaction · 1 assignee · Indexed Sep 16, 2026
document when and how to make ndarray copies of DeviceArrays Open

documentation

jax-ml/jax#390 · 2 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Implement JVP for SVD when full_matrices=True Open

enhancement

jax-ml/jax#508 · 13 comments · 2 reactions · 0 assignees · Indexed Sep 16, 2026
Support autodiff of Eigendecomposition with repeated eigenvalues Open

enhancement

jax-ml/jax#669 · 39 comments · 3 reactions · 0 assignees · Indexed Sep 16, 2026
Jit with unknown shape Open

question

jax-ml/jax#803 · 4 comments · 10 reactions · 1 assignee · Indexed Sep 16, 2026
document memory management (DeviceArrays and where to find them) Open

documentation enhancement

jax-ml/jax#863 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Allow to selectively execute ram-heavy ops on CPU Open

enhancement NVIDIA GPU P2 (eventual)

jax-ml/jax#915 · 2 comments · 5 reactions · 0 assignees · Indexed Sep 16, 2026
Custom tangent types for gradient computations Open

enhancement

jax-ml/jax#918 · 2 comments · 2 reactions · 0 assignees · Indexed Sep 16, 2026
slow pmap allreduce Open

enhancement P2 (eventual) performance

jax-ml/jax#973 · 7 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
feature request: sparse jacobian and sparse hessians Open

enhancement

jax-ml/jax#1032 · 43 comments · 14 reactions · 0 assignees · Indexed Sep 16, 2026
`is` keyword is not preserved under vmap. Open

documentation

jax-ml/jax#1103 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Feature request: extracting intermediate values from custom vjp Open

enhancement

jax-ml/jax#1197 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Clear GPU memory Open

NVIDIA GPU P1 (soon) question

jax-ml/jax#1222 · 22 comments · 6 reactions · 0 assignees · Indexed Sep 16, 2026
Performance Issues with Summing Per-Example Gradients (Caused by Fusions?) Open

bug performance

jax-ml/jax#1232 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Remove the multiple_results field for core.Primitive Open

cleanup

jax-ml/jax#1342 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 16, 2026
jax.numpy.log(2., dtype=np.float32) does not accept dtype arg Open

jax-ml/jax#1364 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Allowing pmap on CPUs by making CPU cores visibles directly on Jax. Open

CPU P1 (soon)

jax-ml/jax#1408 · 17 comments · 29 reactions · 1 assignee · Indexed Sep 16, 2026
Jax for Web? (JS api or web assembly guide) Open

P3 (no schedule)

jax-ml/jax#1472 · 22 comments · 10 reactions · 1 assignee · Indexed Sep 16, 2026
[feature request] Provide access to Pytree to_iterable and from_iterable functions Open

jax-ml/jax#1527 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Differentiable matrix-free linear algebra, optimization and equation solving Open

enhancement NVIDIA GPU P2 (eventual)

jax-ml/jax#1531 · 38 comments · 15 reactions · 0 assignees · Indexed Sep 16, 2026
Limit number of threads when running on CPU Open

documentation

jax-ml/jax#1539 · 29 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Overrides of NumPy functions on JAX arrays Open

enhancement

jax-ml/jax#1565 · 14 comments · 11 reactions · 0 assignees · Indexed Sep 16, 2026
pswapaxes doesn't work with "non-input" tensors Open

jax-ml/jax#1651 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Support for stateful traces Open

jax-ml/jax#1663 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
@differentiable decorator for classes Open

jax-ml/jax#1808 · 6 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Add DCT ops in fftpack Open

enhancement

jax-ml/jax#1839 · 14 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Dynamic sizes in dynamic programming Open

jax-ml/jax#1859 · 2 comments · 0 reactions · 1 assignee · Indexed Sep 16, 2026
Numba bridge Open

contributions welcome NVIDIA GPU P2 (eventual)

jax-ml/jax#1870 · 18 comments · 12 reactions · 1 assignee · Indexed Sep 16, 2026
Let's split the VJP for odeint into JVP and transpose Open

jax-ml/jax#1927 · 6 comments · 11 reactions · 0 assignees · Indexed Sep 16, 2026
float128 please?? Open

enhancement P3 (no schedule)

jax-ml/jax#2007 · 10 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Performance of indexing operations and irregularly sized data Open

NVIDIA GPU performance

jax-ml/jax#2032 · 5 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Function cache for `grad` Open

enhancement

jax-ml/jax#2095 · 7 comments · 1 reaction · 0 assignees · Indexed Sep 16, 2026
scan with gradient checkpointing Open

jax-ml/jax#2139 · 6 comments · 14 reactions · 0 assignees · Indexed Sep 16, 2026
disable_jit should error if jit'd Open

documentation

jax-ml/jax#2148 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Making PRNG output elements consistent for different shapes Open

jax-ml/jax#2155 · 3 comments · 2 reactions · 0 assignees · Indexed Sep 16, 2026
Einsum can be slower that matmul on CPU Open

CPU enhancement P2 (eventual) performance XLA

jax-ml/jax#2160 · 14 comments · 3 reactions · 1 assignee · Indexed Sep 16, 2026
Feature Idea: Tracing Computational cost of algorithms Open

jax-ml/jax#2241 · 2 comments · 5 reactions · 0 assignees · Indexed Sep 16, 2026
Is there a way to store (serialize) expensive first-time use JIT'ted methods? Open

jax-ml/jax#2249 · 6 comments · 3 reactions · 0 assignees · Indexed Sep 16, 2026
[Bug?] Jax shuffle speed versus numpy random shuffle speed Open

performance

jax-ml/jax#2301 · 5 comments · 1 reaction · 1 assignee · Indexed Sep 16, 2026
Gradient of SVD with degenerate singular values becomes NaN Open

P3 (no schedule)

jax-ml/jax#2311 · 14 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Fast JAX compilation when network architecture has per-dataset components. Open

jax-ml/jax#2316 · 2 comments · 1 reaction · 0 assignees · Indexed Sep 16, 2026
jax.jacobian could automatically choose between forward and reverse mode Open

jax-ml/jax#2344 · 2 comments · 1 reaction · 0 assignees · Indexed Sep 16, 2026
[FR] add betaincinv Open

enhancement P3 (no schedule)

jax-ml/jax#2399 · 16 comments · 11 reactions · 0 assignees · Indexed Sep 16, 2026
Outstanding Primitive Rules for Higher Order Automatic Differentiation (jet) Open

jax-ml/jax#2431 · 28 comments · 8 reactions · 0 assignees · Indexed Sep 16, 2026
add attribute to devicearray? Open

jax-ml/jax#2464 · 5 comments · 0 reactions · 0 assignees · Indexed Sep 16, 2026
Please consider adding scipy.special.iv to jax.special Open

contributions welcome enhancement

jax-ml/jax#2466 · 14 comments · 5 reactions · 0 assignees · Indexed Sep 16, 2026
Implement an option for lax.while_loop to specify the maximum number of iterations, to allow reverse differentiation Open

enhancement

jax-ml/jax#2469 · 3 comments · 4 reactions · 1 assignee · Indexed Sep 16, 2026
Implement matrix functions from scipy.linalg Open

contributions welcome

jax-ml/jax#2478 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 16, 2026

jax-ml/jax

Get new issues in your inbox