flax

e36b3f18 - Fix init/apply return annotation