[Flax] Improve Robustness of Back-Prop Tests
This PR makes several modifications to the `test_freeze_feature_encoder` tests for two Flax models, FlaxWav2Vec2 and FlaxSpeechEncoderDecoderModel. This test is in place to verify that feature encoder of the FlaxWav2Vec2 model is correctly frozen when the keyword argument `freeze_feature_encoder` is set to `True`. This is achieved through one dummy forward and backward pass, in which the losses and gradients are compared for the case when the feature encoder is frozen and the case when it is not.

The modifications made include:

- Implementing a check that asserts the outputs of the unfrozen and frozen feature encoder models are _precisely equal_;

- Asserting that the losses of the unfrozen and frozen models are _precisely equal_, instead of _almost equal_ as was previously;

- Asserting that the gradients of the unfrozen feature encoder layers contain at least one non-zero entry, instead of asserting that at least one entry that surpasses a specified tolerance;

- Asserting that the gradients of all unfrozen layers remain _precisely equal_, instead of _almost equal_ as was previously.

The statements that assert unfrozen and frozen gradient values are _precisely equal_ instead of _almost equal_ provides more rigour to the back-propagation test. It verifies that `jax.stop_gradient` has no impact on the forward or backward pass, other than skipping computation of the gradients of any frozen layers. **Crucially**, these changes bypass the need to set a specified tolerance (`tol`) in the `assert_almost_equal` or `assert_difference` functions. Even with seemingly appropriate values for this tolerance (`1e-5`), the tests were prone to failing a small proportion of the time (<2.5%). With this PR, the tests should now pass 100% of the time.
_The documentation is not available anymore as the PR was closed or merged._