test: adjust assertion for my_fn_calls based on optax version in test/test_optimizers.py

[Qazalbash]

We fixed the issue related to optax==0.2.7 in #2137. Now flax has pinned the optax version (google/flax#5225), and it is effective in the latest flax==0.12.4. This PR changes the test case modified in #2137 for optax==0.2.7 and optax<0.2.7 (ref: #2134 (comment))

[fehiepsi]

Maybe be flexible for optax check and removing version checking? That way the test keeps working if optax behavior changes.

[Qazalbash]

Maybe be flexible for optax check and removing version checking? That way the test keeps working if the optax behavior changes.

Do you want a logical statement instead of optax version, like,

assert (my_fn_calls == 1) or (my_fn_calls == 2)

[fehiepsi]

Yes, that would be flexible, maybe calls in [1,2]