fix name scopes to be picked up by quantization config (#294)

* fix name scopes

* Fixed named_scope

* Update name scopes

* Fix name scopes

* Update transformer_wan.py

save memory by deleting params. lint. (#277)

Deletes params pytree after creating TrainState to reduce HBM usage. Lints the code.

The only change is in this line: https://github.com/AI-Hypercomputer/maxdiffusion/compare/deallocate_params_tree?expand=1#diff-69cda939e98b489aca1cc8aa543ecc537d9f4bc6c58fa767cbe01cd3636aacf3R311

version 1