Use select rather than _cond.

PiperOrigin-RevId: 323757367
Change-Id: Id512245c608a6c2709d7748f99308d6502d7a2d9

haiku/_src/moving_averages.py

@@ -57,11 +57,6 @@ def __init__(self, decay, zero_debias=True, warmup_length=0, name=None):
           "average to an initial value. Set zero_debias=False if setting "
           "warmup_length to a non-zero value.")
 
-  def _cond(self, cond, t, f, dtype):
-    """Internal, implements jax.lax.cond without control flow."""
-    c = cond.astype(dtype)
-    return c * t + (1. - c) * f
-
   def initialize(self, value):
     """If uninitialized sets the average to ``zeros_like`` the given value."""
     base.get_state("hidden", value.shape, value.dtype, init=jnp.zeros)
@@ -91,7 +86,7 @@ def __call__(self, value, update_stats=True):
 
     decay = jax.lax.convert_element_type(self._decay, value.dtype)
     if self._warmup_length > 0:
-      decay = self._cond(counter <= 0, 0.0, decay, value.dtype)
+      decay = jax.lax.select(counter <= 0, 0.0, decay)
 
     one = jnp.ones([], value.dtype)
     hidden = base.get_state("hidden", value.shape, value.dtype, init=jnp.zeros)