flax

804a165a - Raise informative error if user tries to naively pass a Module to a jitted function.