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

While2For Raise error Open

EnzymeAD/Enzyme-JAX#1455 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Re-enable special casing for `SliceSimplify` Open

EnzymeAD/Enzyme-JAX#1468 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Batched AD crash for a simple MLP Open

EnzymeAD/Enzyme-JAX#1469 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Conv with splatted constant LHS/RHS simplification Open

EnzymeAD/Enzyme-JAX#1474 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Extend DotGeneralSimplify to handle Broadcasted Scalars Open

good first issue

EnzymeAD/Enzyme-JAX#1475 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Run downstream Reactant tests Open

EnzymeAD/Enzyme-JAX#1477 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
transpose(A * A) where A is symmetric Open

EnzymeAD/Enzyme-JAX#1481 · 7 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
`whileOp` reduction into `reduce` Open

EnzymeAD/Enzyme-JAX#1506 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Remove cudnn debug xla override commit Open

EnzymeAD/Enzyme-JAX#1534 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
busted optimization Open

EnzymeAD/Enzyme-JAX#1537 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Raise to stablehlo error Open

EnzymeAD/Enzyme-JAX#1554 · 5 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Missing LICM Passes Open

EnzymeAD/Enzyme-JAX#1558 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
While DUS to Broadcast In Dim Open

EnzymeAD/Enzyme-JAX#1563 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
convolution with batch_group_count -> regular convolution Open

EnzymeAD/Enzyme-JAX#1572 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Symmetric Matrix Checks Open

EnzymeAD/Enzyme-JAX#1609 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Dot General Diagonal Simplifications Open

EnzymeAD/Enzyme-JAX#1617 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Slice of While(DUS) along iteration dimension Open

EnzymeAD/Enzyme-JAX#1636 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Reverse of constant -> constant Open

EnzymeAD/Enzyme-JAX#1646 · 0 comments · 1 reaction · 1 assignee · Indexed Sep 15, 2026
Slice of reverse -> reverse of slice Open

EnzymeAD/Enzyme-JAX#1647 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Reverse(reverse) -> noop Open

EnzymeAD/Enzyme-JAX#1648 · 0 comments · 1 reaction · 1 assignee · Indexed Sep 15, 2026
filled syrk -> syrk if only specific triangle accessed Open

EnzymeAD/Enzyme-JAX#1652 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Support for automatic differentiation of JAX code Open

EnzymeAD/Enzyme-JAX#1661 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Miscompiles for VJP when using jax.lax.associative_scan Open

EnzymeAD/Enzyme-JAX#1664 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Cannot raise while with bounds checks Open

raising

EnzymeAD/Enzyme-JAX#1672 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
`ERROR: LoadError: MethodError: no method matching Float64(::Reactant.TracedRNumber{Float64})` for run with ImplicitFreeSurface on Oceananigans Open

EnzymeAD/Enzyme-JAX#1674 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Check if dimension is dynamic and use dynamic_reshape instead of reshape in affinetostablehlo raising Open

EnzymeAD/Enzyme-JAX#1678 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Add enzymexla.free op [and convert cudaFree into it] Open

EnzymeAD/Enzyme-JAX#1679 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Explicitly track all possible index access in stablehlo.while Open

EnzymeAD/Enzyme-JAX#1681 · 0 comments · 0 reactions · 1 assignee · Indexed Sep 15, 2026
Optimize Gather/Scatter into reshape with slice Open

EnzymeAD/Enzyme-JAX#1682 · 5 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
`Concat(Slice(),...)` fuse into `Gather`? Open

EnzymeAD/Enzyme-JAX#1688 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Element type Conversion and offset Open

EnzymeAD/Enzyme-JAX#1708 · 0 comments · 1 reaction · 1 assignee · Indexed Sep 15, 2026
BatchOp(BroadcastInDim(), ....) -> Broadcast(BatchOp) Open

EnzymeAD/Enzyme-JAX#1720 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Expand DotGeneral CSE Pass to handle more equivalence Open

EnzymeAD/Enzyme-JAX#1725 · 0 comments · 1 reaction · 1 assignee · Indexed Sep 15, 2026
DS(DUS) with non-intersecting update Open

EnzymeAD/Enzyme-JAX#1731 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Negate of constant, constant propagation Open

EnzymeAD/Enzyme-JAX#1732 · 0 comments · 1 reaction · 1 assignee · Indexed Sep 15, 2026
Remaining all-reduce Open

EnzymeAD/Enzyme-JAX#1740 · 8 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Select or DUS-based version of ConcatToPad Open

EnzymeAD/Enzyme-JAX#1741 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Add Revolve checkpointing Open

EnzymeAD/Enzyme-JAX#1743 · 4 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
GB25 CI with > halo^2 number of devices (e.g 128) Open

EnzymeAD/Enzyme-JAX#1751 · 0 comments · 0 reactions · 1 assignee · Indexed Sep 15, 2026
Index analysis of while dynamic slice Open

EnzymeAD/Enzyme-JAX#1752 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Shardy all-to-all error Open

EnzymeAD/Enzyme-JAX#1761 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
DUS of DUS -> DUS of extend Open

EnzymeAD/Enzyme-JAX#1767 · 2 comments · 1 reaction · 1 assignee · Indexed Sep 15, 2026
Remaining all-reduce part 2 Open

EnzymeAD/Enzyme-JAX#1774 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
reduce of broadcast Open

EnzymeAD/Enzyme-JAX#1788 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
hlo opt gives incorrect results for neural gcm Open

python

EnzymeAD/Enzyme-JAX#1796 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
re-enable Keras benchmarks Open

python

EnzymeAD/Enzyme-JAX#1797 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
re-enable maxtext benchmark Open

python

EnzymeAD/Enzyme-JAX#1798 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Enable mlir ad for existing python tests Open

python

EnzymeAD/Enzyme-JAX#1799 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Add documentation on how to actually use with jax Open

python

EnzymeAD/Enzyme-JAX#1800 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
long chain of self additions to mul Open

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

EnzymeAD/Enzyme-JAX

Get new issues in your inbox