[JAX] Use default factory for not sharing mutable default values

[zlsh80826]

Description

This PR addresses an issue caused by using mutable default values in dataclass fields, which can lead to unexpected behavior. More details are available in the official documentation.

Starting with Python 3.11, an explicit error is raised when mutable values are used as default values in dataclass fields. To resolve this, the default_factory parameter is now used as recommended in the documentation. This approach ensures that new instances of mutable types are created for each dataclass instance, avoiding shared state and related issues.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refractor

Changes

Using default_factory to set an non-frozen dataclass as the default value

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• [] I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[zlsh80826]

/te-ci jax L1

[phu0ngng]

The CI failed as follows:

• on A100 node withtest_model_parallel_encoder.py::TestEncoder::test_te_bf16
• L1 distributed test with test_distributed_layernorm.py

I think these failures are unrelated to the changes introduced in this PR.
Therefore, we can merge this PR to unblock the CI for now while looking into those failures.