flax
add graph_updates argument to jit
#5317
Merged

Loading