Replaced jax.lax.select with jnp.where

jmp/_src/loss_scale.py

@@ -143,12 +143,12 @@ def adjust(self, grads_finite: jnp.ndarray) -> "DynamicLossScale":
     """Returns the next state dependent on whether grads are finite."""
     assert grads_finite.ndim == 0, "Expected boolean scalar"
 
-    first_finite = lambda a, b: jax.lax.select(jnp.isfinite(a).all(), a, b)
-    loss_scale = jax.lax.select(
+    first_finite = lambda a, b: jnp.where(jnp.isfinite(a).all(), a, b)
+    loss_scale = jnp.where(
         grads_finite,
 
         # When grads are finite increase loss scale periodically.
-        jax.lax.select(
+        jnp.where(
             self.counter == (self.period - 1),
             first_finite(self.loss_scale * self.factor,
                          self.loss_scale),
@@ -188,4 +188,4 @@ 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_map(functools.partial(jnp.where, pred), a, b)