EnzymeAD/Enzyme-JAX

View on GitHub

Custom Bindings for Enzyme Automatic Differentiation Tool and Interfacing with JAX.

Stars: 131
Forks: 53
Open beginner issues: 0
Indexed issues: 291
Avg merge: 1d 10h
Merged PRs (30d): 193

Dominant language: MLIR
License: No license data
Last GitHub push: Sep 14, 2026
Latest indexed: Sep 15, 2026
Contributing guide: No contributing guide
Code of conduct: No code of conduct
Beginner labels: good first issue help wanted

Open All

291 open issues indexed

Add support for `jacrev`, `jacfwd`, `hessian`, `vmap` Open

EnzymeAD/Enzyme-JAX#9 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
HLO Canonicalizations Todo list Open

EnzymeAD/Enzyme-JAX#54 · 31 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Tracking issue for missing HLO derivatives Open

good first issue help wanted

EnzymeAD/Enzyme-JAX#88 · 4 comments · 0 reactions · 1 assignee · Indexed Sep 15, 2026
GPU CI of Jax not thinking its running on GPU Open

EnzymeAD/Enzyme-JAX#122 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Profiling API's Open

EnzymeAD/Enzyme-JAX#129 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Export issue with missing BLAS symbol Open

EnzymeAD/Enzyme-JAX#135 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Local build is broken on macOS Open

EnzymeAD/Enzyme-JAX#146 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Tracking issue for missing Batch Op Interface Open

good first issue help wanted

EnzymeAD/Enzyme-JAX#152 · 7 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
TupleOp Open

EnzymeAD/Enzyme-JAX#157 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Constprop through chlo.lgamma Open

EnzymeAD/Enzyme-JAX#179 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
An enzyme optimization benchmark Open

EnzymeAD/Enzyme-JAX#180 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Optimization Passes for `dynamic_gather` Open

EnzymeAD/Enzyme-JAX#184 · 2 comments · 1 reaction · 1 assignee · Indexed Sep 15, 2026
Missing Transpose Optimizations Open

EnzymeAD/Enzyme-JAX#217 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Kernel SROA of arguments Open

EnzymeAD/Enzyme-JAX#227 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Incorrect IR generated for Vector Mode AD Open

EnzymeAD/Enzyme-JAX#235 · 4 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Kernel Optimizations Open

EnzymeAD/Enzyme-JAX#238 · 1 comment · 0 reactions · 1 assignee · Indexed Sep 15, 2026
`slice(transpose(...))` --> `transpose(slice(...))` if transpose is only being used in slice Open

EnzymeAD/Enzyme-JAX#257 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Raise NVVM threadidx to GPU threadidx [and similar] Open

EnzymeAD/Enzyme-JAX#274 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
extui and related arithmetic instructions not eliminated Open

EnzymeAD/Enzyme-JAX#279 · 3 comments · 0 reactions · 1 assignee · Indexed Sep 15, 2026
Scatter Optimizations with Window Dims specified Open

EnzymeAD/Enzyme-JAX#308 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Propagate Constants into Body of Operations Open

EnzymeAD/Enzyme-JAX#317 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
More optimization patterns as required Open

EnzymeAD/Enzyme-JAX#370 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Affine if simplification Open

EnzymeAD/Enzyme-JAX#392 · 8 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
GB Stablehlo optimization Open

EnzymeAD/Enzyme-JAX#402 · 7 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Add support to raise an affine.for to stablehlo.while Open

EnzymeAD/Enzyme-JAX#442 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Pull repeated ops from conditional branches Open

EnzymeAD/Enzyme-JAX#444 · 4 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Illegal frem error Open

EnzymeAD/Enzyme-JAX#482 · 0 comments · 1 reaction · 0 assignees · Indexed Sep 15, 2026
CSE of not Open

EnzymeAD/Enzyme-JAX#490 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Transpose of a while Open

EnzymeAD/Enzyme-JAX#506 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
BatchNorm Inference fusion with Dense or Conv Open

EnzymeAD/Enzyme-JAX#533 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
IR right before stablehlo raise Open

EnzymeAD/Enzyme-JAX#536 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Remaining memref.load Open

EnzymeAD/Enzyme-JAX#543 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Test cases for improved for loop raising Open

EnzymeAD/Enzyme-JAX#545 · 7 comments · 0 reactions · 1 assignee · Indexed Sep 15, 2026
concat reshape simplification Open

EnzymeAD/Enzyme-JAX#557 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Subtract reduction raising Open

EnzymeAD/Enzyme-JAX#560 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Further generalize reductions Open

EnzymeAD/Enzyme-JAX#566 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Transpose Optimizations Missing (Move Transposes Up) Open

EnzymeAD/Enzyme-JAX#568 · 5 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Push reshapes up Open

EnzymeAD/Enzyme-JAX#575 · 29 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
ERROR: LoadError: "failed to run pass manager on module" on Reactant script with Oceananigans Open

EnzymeAD/Enzyme-JAX#578 · 10 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
DUS creates Allgather operations not removed by XLA Open

EnzymeAD/Enzyme-JAX#596 · 10 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Simplify/eliminate DUS Open

EnzymeAD/Enzyme-JAX#605 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Failed raising Open

EnzymeAD/Enzyme-JAX#614 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Padding remains in baroclinic model Open

EnzymeAD/Enzyme-JAX#622 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
`DEADLINE_EXCEEDED: GetKeyValue() timed out` Open

EnzymeAD/Enzyme-JAX#623 · 5 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Slice access analysis to eliminate slices and concats inside loops Open

EnzymeAD/Enzyme-JAX#625 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Unused loop iterarg returning constant Open

EnzymeAD/Enzyme-JAX#627 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Slicing to conv Open

EnzymeAD/Enzyme-JAX#628 · 9 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Concat of operation with slices on boundary -> dynamic_slice Open

EnzymeAD/Enzyme-JAX#630 · 7 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Conv creates allgather Open

EnzymeAD/Enzyme-JAX#653 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
OOM for large enough models Open

EnzymeAD/Enzyme-JAX#671 · 8 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026

EnzymeAD/Enzyme-JAX

Get new issues in your inbox