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

Improve concat like comms Open

EnzymeAD/Enzyme-JAX#1844 · 0 comments · 1 reaction · 2 assignees · Indexed Sep 15, 2026
raise scan-like operations Open

EnzymeAD/Enzyme-JAX#1850 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
XLA-level optimization Open

EnzymeAD/Enzyme-JAX#1851 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Optimize to extend of rotate: Open

EnzymeAD/Enzyme-JAX#1854 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Extend To DUS Open

EnzymeAD/Enzyme-JAX#1855 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
More wrap rotate fixups Open

EnzymeAD/Enzyme-JAX#1858 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
reduce_window fusions Open

EnzymeAD/Enzyme-JAX#1859 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
factor out the common matvec Open

EnzymeAD/Enzyme-JAX#1863 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
[polybench] kernel symm raising failure Open

EnzymeAD/Enzyme-JAX#1864 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
[polybench] blas trmm raising failure Open

EnzymeAD/Enzyme-JAX#1866 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
dot_general reshape of non-contracting / non-batching dims Open

EnzymeAD/Enzyme-JAX#1867 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
XLA patch, flip select true/false in mask Open

EnzymeAD/Enzyme-JAX#1879 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Broadcast transpose reshape fixup Open

EnzymeAD/Enzyme-JAX#1881 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Communication: Extend(Extend(x)) Open

EnzymeAD/Enzyme-JAX#1882 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
More aggressive slice of update_without_corners Open

EnzymeAD/Enzyme-JAX#1883 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
`extendrotate` seems broken on mac Open

EnzymeAD/Enzyme-JAX#1889 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
more scatter-gather patterns Open

EnzymeAD/Enzyme-JAX#1891 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
fuse `scatter(in, ind, elem_op(gather(in, ind), update))` into scatter Open

EnzymeAD/Enzyme-JAX#1892 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Add an IR level option to disable auto-batching for a specific while loop Open

EnzymeAD/Enzyme-JAX#1893 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Fix looped fill halo region to fully remove loop Open

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

EnzymeAD/Enzyme-JAX#1919 · 6 comments · 1 reaction · 2 assignees · Indexed Sep 15, 2026
slice(gather(x, ind)) -> gather(x, slice(ind)) Open

EnzymeAD/Enzyme-JAX#1924 · 3 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
DUS of select of slice [which is noop to dus], and actual val -> select(oldval, DUS(...)) Open

EnzymeAD/Enzyme-JAX#1926 · 0 comments · 1 reaction · 2 assignees · Indexed Sep 15, 2026
Batching of BatchNorm* Ops is Busted Open

EnzymeAD/Enzyme-JAX#1947 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Optimize fft(0) to be 0 Open

EnzymeAD/Enzyme-JAX#1951 · 2 comments · 1 reaction · 2 assignees · Indexed Sep 15, 2026
MultiRotate Op and MultiSliceOp Open

EnzymeAD/Enzyme-JAX#1949 · 4 comments · 1 reaction · 2 assignees · Indexed Sep 15, 2026
scatter to slice optimization missing Open

EnzymeAD/Enzyme-JAX#1957 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
FFT(aA) where a is a scalar Open

EnzymeAD/Enzyme-JAX#1987 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Decompose elementwise multiply for complex numbers (one has zero imag) Open

complex

EnzymeAD/Enzyme-JAX#1988 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
imag, real, abs of `exp(complex number)` Open

complex

EnzymeAD/Enzyme-JAX#1990 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
`lowering_gpu-recognize2.mlir.test` crashes in debug build Open

EnzymeAD/Enzyme-JAX#2004 · 2 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Hoise Enzyme AD Operations from While Loop by specifying `width` Open

auto-batching

EnzymeAD/Enzyme-JAX#2027 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Not removed reshape Open

EnzymeAD/Enzyme-JAX#2031 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Two broadcast optimizations Open

EnzymeAD/Enzyme-JAX#2038 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Missing v0.0.11 pypi wheels Open

EnzymeAD/Enzyme-JAX#2055 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Rewrite conditional binary op with binary op with init value Open

EnzymeAD/Enzyme-JAX#2057 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
concatenate slice optimization Open

EnzymeAD/Enzyme-JAX#2087 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Multi rotate cse Open

EnzymeAD/Enzyme-JAX#2090 · 6 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
TPU-Blocker: custom call memory operations in f64 being rewritten Open

EnzymeAD/Enzyme-JAX#2126 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
missing reduce constant prop Open

EnzymeAD/Enzyme-JAX#2141 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Optimize transposed operands / transpose(result) for custom_call Open

EnzymeAD/Enzyme-JAX#2172 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
No proper stack trace for enzymexla-mlir-opt Open

bazel

EnzymeAD/Enzyme-JAX#2185 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
LU(Symmetric) \ N -> Bunch Kaufman (sytrf) Open

EnzymeAD/Enzyme-JAX#2186 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Issues under raising opaque SSA-defined bounds to ISL Open

EnzymeAD/Enzyme-JAX#2212 · 4 comments · 0 reactions · 4 assignees · Indexed Sep 15, 2026
Redirect Slice through UpdateWithoutBorders Open

EnzymeAD/Enzyme-JAX#2265 · 0 comments · 0 reactions · 0 assignees · Indexed Sep 15, 2026
Pad alignment for TPU perf Open

EnzymeAD/Enzyme-JAX#2305 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
C++ raise to tpu : Failed to infer result type(s): Open

EnzymeAD/Enzyme-JAX#2331 · 0 comments · 1 reaction · 0 assignees · Indexed Sep 15, 2026
ConvertParallelToGPU1 pass crashes Open

help wanted

EnzymeAD/Enzyme-JAX#2358 · 1 comment · 0 reactions · 3 assignees · Indexed Sep 15, 2026
Decouple from JAX Open

EnzymeAD/Enzyme-JAX#2375 · 1 comment · 0 reactions · 0 assignees · Indexed Sep 15, 2026
autodiff through while loops with state-dependent break Open

EnzymeAD/Enzyme-JAX#2565 · 0 comments · 0 reactions · 1 assignee · Indexed Sep 15, 2026

EnzymeAD/Enzyme-JAX

Get new issues in your inbox