Remove two warnings that were showing when using jmp.

Fixes #39.

PiperOrigin-RevId: 505639478

Author: tomhennigan

jmp/_src/loss_scale.py

@@ -105,7 +105,7 @@ class DynamicLossScale:
 
   Typical usage of this class will be something like:
 
-  >>> loss_scale = jmp.DynamicLossScale(jnp.asarray(2 ** 15))
+  >>> loss_scale = jmp.DynamicLossScale(jnp.asarray(2. ** 15))
   >>> for _ in range(num_steps):
   ...   # compute loss
   ...   loss = loss_scale.scale(loss)
@@ -121,7 +121,7 @@ class DynamicLossScale:
   period: int = 2000
   factor: int = 2
   min_loss_scale: jnp.ndarray = dataclasses.field(
-      default_factory=lambda: np.ones([], np.int32))
+      default_factory=lambda: np.ones([], np.float32))
 
   def __post_init__(self) -> None:
     warn_if_not_floating(self.loss_scale, "loss_scale")
@@ -183,7 +183,7 @@ def adjust(self, grads_finite: jnp.ndarray) -> "DynamicLossScale":
 
 def all_finite(tree) -> jnp.ndarray:
   """Returns a scalar ndarray indicating whether the input arrays are finite."""
-  leaves = jax.tree_leaves(tree)
+  leaves = jax.tree_util.tree_leaves(tree)
   if not leaves:
     return jnp.array(True)
   else:
@@ -195,7 +195,7 @@ def all_finite(tree) -> jnp.ndarray:
 def select_tree(pred: jnp.ndarray, a: T, b: T) -> T:
   """Selects a pytree based on the given predicate."""
   assert pred.ndim == 0 and pred.dtype == jnp.bool_, "expected boolean scalar"
-  return jax.tree_map(functools.partial(jax.lax.select, pred), a, b)
+  return jax.tree_util.tree_map(functools.partial(jax.lax.select, pred), a, b)
 
 
 def warn_if_not_floating(x: Union[jnp.ndarray, object], var_name: str) -> None:

jmp/_src/loss_scale_test.py

@@ -14,6 +14,8 @@
 # ==============================================================================
 """Tests for jmp._src.loss_scale."""
 
+import warnings
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -51,15 +53,20 @@ def test_static_loss_scale(self, cls, scale):
   )
   def test_static_empty_trees(self, create):
     loss_scale = create()
-    self.assertEmpty(jax.tree_leaves(loss_scale))
+    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)
+    self.assertEmpty(logged_warnings)
 
   def test_dynamic_loss_scale_tree(self):
     scale = jnp.ones([])
     counter = jnp.zeros([], jnp.int32)
     period = 2000
     factor = 2
     loss_scale = jmp.DynamicLossScale(scale, counter, period, factor)
-    self.assertEqual(jax.tree_leaves(loss_scale), [scale, counter])
+    self.assertEqual(jax.tree_util.tree_leaves(loss_scale), [scale, counter])
     self.assertEqual(jax.tree_util.tree_map(lambda x: x, loss_scale),
                      loss_scale)
 

jmp/_src/policy_test.py

@@ -54,8 +54,8 @@ def skip_if_unsupported(dtype):
 class PolicyTest(parameterized.TestCase):
 
   def assert_dtypes_equal(self, tree_a, tree_b):
-    jax.tree_map(lambda a, b: self.assertEqual(a.dtype, b.dtype), tree_a,
-                 tree_b)
+    jax.tree_util.tree_map(lambda a, b: self.assertEqual(a.dtype, b.dtype),
+                           tree_a, tree_b)
 
   @parameterized.parameters(*it.product(DTYPES, NUMPYS))
   def test_policy_cast_to_param(self, dtype, np_):