flax

[nnx] fix Variable overloads and add shape/dtype properties

#4049

Merged