[Flax tests] fix test_model_outputs_equivalence
# What does this PR do?

The `test_model_outputs_equivalence` common test is never run because `recursive_check` call is at the wrong place. This PR calls `recursive_check` in `check_equivalence` and also fixes the `set_nan_tensor_to_zero` by using an equivalent jax function since jax arrays do not support index assignment.
_The documentation is not available anymore as the PR was closed or merged._<|||||>Thanks a lot for fixing it