flax

fix eval_shape's _to_variable

#5316

Merged