Add pickle/cloudpickle round-trip tests for Linen modules (#1481)

[kapellirohith]

Fixes #1481

Description

This PR introduces unit tests ensuring that Flax Linen Modules can be safely
serialized and deserialized via pickle and cloudpickle without leaking
internal binding state. All new assertions are isolated within
tests/pickle_test.py; no library code is modified.

Changes Introduced

• Created ModulePickleTest with a shared helper assert_roundtrip_identical
  that initializes the module before dumping (so it has been bound at least
  once), verifies the restored module comes back unbound (scope is None),
  and asserts apply outputs and freshly-init-ed parameters are
  bit-identical before and after round-tripping.
• Added 6 test cases covering:
  • nn.Dense with explicitly picklable initializers via stdlib pickle
    (with full object equality), and default nn.Dense via cloudpickle.
  • A compact MLP with nested submodules and random default initializers,
    via both picklers.
  • A module class defined in a local scope (the notebook use case fixed in
    Fix cloudpickle serialization of wrapped Module methods. #1475), via cloudpickle.

Technical Findings

• JAX initializer closure limitation: default nn.Dense fails stdlib
  pickling because the default kernel initializer (lecun_normal) is a
  closure created inside jax.nn.initializers.variance_scaling, which pickle
  cannot serialize by reference. Rather than modifying core code, this is
  documented via an expected-failure canary test
  (test_pickle_default_initializer_limitation) — if it ever starts failing,
  the limitation was lifted and the tests should be updated.
• Cloudpickle function identity: cloudpickle restores initializer closures
  by value, so restored function fields are equivalent copies rather than
  identical objects. The default-Dense cloudpickle test therefore validates
  behavioral equality (outputs and re-initialized parameters) instead of
  dataclass equality.
• dill interplay (raised in the issue): nothing in this test file's import
  chain pulls in dill, so the dill/cloudpickle test-setup conflict described
  in Test pickle support #1481 does not affect these tests. cloudpickle>=3.0.0 is already in the
  testing extras, so CI needs no dependency changes.

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs.
• This change is discussed in a Github issue/discussion: Test pickle support #1481
• The documentation and docstrings adhere to the documentation guidelines.
• This change includes necessary high-coverage tests. (No quality testing = no merge!)

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.