Fix JAX transformation compatibility with functional updates#2

Merged

juehang merged 2 commits into

mainfrom

fix/jax-transformation-compatibility

Nov 20, 2025

Commits on Nov 20, 2025

test(dataloader): add tests for JAX grad compatibility in transformation methods
juehang
committed
refactor(dataloader): improve JAX transformation compatibility with functional updates
juehang
committed