flax

Change the type signature of nn.compact to transparently pass its input through

#1041

Merged