Update the code to not use soon-to-be-deprecated Chex assertions.

PiperOrigin-RevId: 551473625
google-deepmind · Jul 27, 2023 · 1afd608 · 1afd608
1 parent cfa6f78
commit 1afd608
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optax/_src/wrappers_test.py b/optax/_src/wrappers_test.py
@@ -519,7 +519,7 @@ def masked_negate(updates):
 
     with self.subTest('tree_map_params'):
       result = state_utils.tree_map_params(init_fn, lambda v: v, state)
-      chex.assert_tree_all_equal_structs(result, state)
+      chex.assert_trees_all_equal_structs(result, state)
 
     updates, state = update_fn(input_updates, state, params)
     chex.assert_trees_all_close(updates, correct_updates)