rlx-autodiff 0.2.2

JAX-shaped transforms for RLX MIR — autodiff, JVP/HVP, vmap

Coverage
83.87%
26 out of 31 items documented0 out of 20 items with examples
Size
Source code size: 265.56 kB This is the summed size of all the files inside the crates.io package for this release.
Documentation size: 1.25 MB This is the summed size of all files generated by rustdoc for all configured targets
Ø build duration
this release: 8s Average build duration of successful builds.
all releases: 8s Average build duration of successful builds in releases after 2024-10-23.
Links
Homepage
MIT-RLX/rlx
3 0 0
crates.io
Dependencies
- rlx-fusion ^0.2.1 normal
- rlx-ir ^0.2.2 normal
Versions
Owners

rlx-autodiff

JAX-shaped transforms on RLX MIR: reverse-mode grad_with_loss, forward-mode jvp / hvp, and vmap (leading-axis batching).

Depends on rlx-ir and [rlx-fusion] (unfuse fused ops before AD when needed).

What's here

autodiff — reverse-mode AD; fused-op VJPs, control flow (If / While / Scan), custom-fn inlining for AD.
autodiff_fwd — forward-mode AD.
prepare_ad — prepare_graph_for_ad, MIR/module preparation.
vmap — batched function transform.
legalize_reduce — reduce legalization helpers for training graphs.
FFT AD — VJP/JVP rules for Op::Fft (unitary / norm-aware).

Feature

Enable via rlx-opt with feature training (default), or depend on this crate directly for a minimal AD-only dep tree.

Build / test

cargo test -p rlx-autodiff

License

GPL-3.0-only.