test: adjust assertion for my_fn_calls based on optax version in test/test_optimizers.py (#2140)

* Revert "test: update test effected by the `optax==0.2.7` release (#2137)"

This reverts commit 7d67ff5.

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

* fix: remove test skip

* fix: remove optax version checks from the test

Author: Qazalbash

test/test_optimizers.py

@@ -123,4 +123,8 @@ def my_fn(state, g):
     state = my_fn(state, jnp.ones(10) * 2.0)
     state = my_fn(state, jnp.ones(10) * 3.0)
 
-    assert my_fn_calls == 1
+    if uses_value_arg:
+        # Dtype is different on the first call vs the rest of the calls
+        assert my_fn_calls in (1, 2)
+    else:
+        assert my_fn_calls == 1