flax

Generalize attribute assignment to support passing module instances

#976

Merged