code: def test_warmup(self):
      ema = moving_averages.ExponentialMovingAverage(
          0.5, warmup_length=2, zero_debias=False
      )
      random_input = jax.random.uniform(
          jax.random.PRNGKey(428), shape=(2, 3, 4)
      )
  
      # The ema should be equal to the input for the first two calls.
      np.testing.assert_allclose(random_input[0], ema(random_input[0]))
      np.testing.assert_allclose(random_input[0], ema(random_input[0]))
  
      # After the warmup period, with decay = 0.5 it should be halfway between the
      # first two inputs and the new input.
      np.testing.assert_allclose(
          (random_input[0] + random_input[1]) / 2, ema(random_input[1])
      )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import basic
  from haiku._src import moving_averages
  from haiku._src import multi_transform
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
  import jax.random
  import numpy
  import tree
label_desc: Use image node_id 2 for calling the MovingAveragesTest obj's underlying member method code with example usage: obj.test_warmup() without return types