Fix or ignore some pytype errors related to jnp.ndarray == jax.Array.

PiperOrigin-RevId: 511782544
google-deepmind · Feb 23, 2023 · 785390a · 785390a
1 parent 52a66ae
commit 785390a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jmp/_src/loss_scale_test.py b/jmp/_src/loss_scale_test.py
@@ -49,15 +49,15 @@ def test_static_loss_scale(self, cls, scale):
 
   @parameterized.named_parameters(
       ("NoOpLossScale", jmp.NoOpLossScale),
-      ("StaticLossScale", lambda: jmp.StaticLossScale(0)),
+      ("StaticLossScale", lambda: jmp.StaticLossScale(0)),  # pytype: disable=wrong-arg-types  # jax-ndarray
   )
   def test_static_empty_trees(self, create):
     loss_scale = create()
     self.assertEmpty(jax.tree_util.tree_leaves(loss_scale))
 
   def test_dynamic_loss_scale_no_warnings(self):
     with warnings.catch_warnings(record=True) as logged_warnings:
-      jmp.DynamicLossScale(2. ** 15)
+      jmp.DynamicLossScale(2. ** 15)  # pytype: disable=wrong-arg-types  # jax-ndarray
     self.assertEmpty(logged_warnings)
 
   def test_dynamic_loss_scale_tree(self):