flax

[nnx] Add complex test with scan + batchnorm + dropout

#3567

Merged