Avoid scalar conversion of non-scalar arrays

Scalar conversion of non-scalar arrays has been deprecated and has raised a warning since JAX v0.4.16, and will soon result in an error (see google/jax#19181).

PiperOrigin-RevId: 595815388

optax/_src/wrappers_test.py

@@ -175,18 +175,18 @@ def fn_update(params, opt_state, x):
     for step in range(2):
       params, opt_state = fn_update(params, opt_state, two)
       self.assertFalse(bool(opt_state.last_finite))
-      self.assertEqual(step + 1, int(opt_state.notfinite_count))
+      self.assertEqual(step + 1, opt_state.notfinite_count.item())
     # Next successful param update
     params, opt_state = fn_update(params, opt_state, half)
     self.assertTrue(bool(opt_state.last_finite))
     # Again 2 rejected param updates
     for step in range(2):
       params, opt_state = fn_update(params, opt_state, two)
       self.assertFalse(bool(opt_state.last_finite))
-      self.assertEqual(step + 1, int(opt_state.notfinite_count))
+      self.assertEqual(step + 1, opt_state.notfinite_count.item())
     # Next param update with NaN is accepted since we reached maximum
     _, opt_state = fn_update(params, opt_state, two)
-    self.assertEqual(5, int(opt_state.total_notfinite))
+    self.assertEqual(5, opt_state.total_notfinite.item())
 
   @chex.variants(with_jit=True, without_jit=True, with_pmap=True)
   def test_multi_steps(self):