flax

bc689087 - Perform shape checks for self.param AFTER unboxing, make shmap users happy