flax

205ed408 - add support for hijax Variables in nnx transforms