[JAX] Adapt to Flax 0.7.1

[mingxu1067]

• Cast Flax collections to FrozenDict as WAR to adapt Flax 0.7.1.
• Followed migrate to the normal dict to apapt to Flax 0.7.1.
  The root cause is that FrozenDict is removed from Flax 0.7.1, instead module.init() returns python built-in dict which has different behavior of pop with FrozenDict. FrozenDict's pop reuturns two outputs, remaining FrozenDict and popped value, but python built-in dict only returns popped value then in-place delete it in orginal object.

Meanwhile, add fully remove FrozenDict into TODO and will schedule time to fix internally.

Refer to #352

[mingxu1067]

@zlsh80826, @jeng1220, @nouiz for viz.

[mingxu1067]

/te-ci

[zlsh80826]

Hi Ming,
There is a guide on how to migrate to the normal dict. I think this would be better, the changes are small.

 def generate_layer(layer_cls, init_rng, diff_inputs, no_diff_inputs):
layer = layer_cls()
variables = layer.init(init_rng, *diff_inputs, *no_diff_inputs)
- others, params = variables.pop('params')
+ others, params = flax.core.pop(variables, 'params')
del variables
return layer, params, others
@@ -126,7 +126,7 @@ class TestEncoderLayer:
def sync_params(ref, target, attrs):
fuse_qkv = attrs.get(_KEY_OF_FUSE_QKV_PARAMS, True)
- unfreeze_target = target.unfreeze()
+ unfreeze_target = flax.core.unfreeze(target)
if fuse_qkv:
unfreeze_target['attention']['qkv']['kernel'] = \
jnp.reshape(ref['attention']['qkv']['kernel'],

[mingxu1067]

Thanks Reese for pointing this out, apopt this changes.

[mingxu1067]

/te-ci

[nouiz]

Need to update the min flax version somewhere.

[mingxu1067]

Need to update the min flax version somewhere.

Done, by adding it to requirements.txt in examples/jax.

[mingxu1067]

/te-ci

[mingxu1067]

/te-ci

[yhtang]

Confirming that TE's JAX tests in JAX-Toolbox is fixed now:

• https://github.com/NVIDIA/JAX-Toolbox/actions/runs/5746725690/job/15577049028
• https://github.com/NVIDIA/JAX-Toolbox/actions/runs/5746725690/job/15577048581
• https://github.com/NVIDIA/JAX-Toolbox/actions/runs/5746725690/job/15577048688
• https://github.com/NVIDIA/JAX-Toolbox/actions/runs/5746725690/job/15577048798
• https://github.com/NVIDIA/JAX-Toolbox/actions/runs/5746725690/job/15577048896

[nouiz]

Any idea why it is fixed? Should we drop this PR?

[timmoon10]

It looks like JAX-Toolbox is fetching this branch, so we should merge soon.

We should also update the Flax version in setup.py:

TransformerEngine/setup.py

Line 298 in 85928d0

add_unique(install_reqs, ["jax", "flax"])