Update the code to not use the deprecated Chex assertions.

PiperOrigin-RevId: 549693785
google-deepmind · Jul 20, 2023 · 2a8a517 · 2a8a517
1 parent 815da81
commit 2a8a517
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/optax/_src/factorized_test.py b/optax/_src/factorized_test.py
@@ -43,7 +43,7 @@ def test_scale_by_factored_rms(self):
 
     updates, state = transform_fn(self.per_step_updates, state, params)
     chex.assert_tree_all_finite((params, updates, state))
-    chex.assert_tree_all_equal_shapes(params, updates)
+    chex.assert_trees_all_equal_shapes(params, updates)
 
 
 if __name__ == '__main__':

diff --git a/optax/_src/wrappers_test.py b/optax/_src/wrappers_test.py
@@ -569,7 +569,7 @@ def test_masked_state_structure(self):
         'a': [jnp.zeros(1), (jnp.zeros(2), wrappers.MaskedNode())],
         'b': wrappers.MaskedNode()
     }
-    chex.assert_tree_all_equal_structs(trace, expected_trace)
+    chex.assert_trees_all_equal_structs(trace, expected_trace)
 
   def test_masked_state_is_compatible_with_deepmind_tree(self):
     """Checks that the masked state is compatible with deepmind/tree.