code: def test_tree_update_stats(self):
      def f():
          return basic.Linear(output_size=2, b_init=jnp.ones)(
              jnp.zeros([6])
          )
  
      init_fn, _ = transform.transform(f)
      params = init_fn(random.PRNGKey(428))
  
      def g(x):
          """This should never update internal stats."""
          return moving_averages.EMAParamsTree(0.2)(
              x, update_stats=False
          )
  
      init_fn, apply_fn_g = multi_transform.without_apply_rng(
          transform.transform_with_state(g)
      )
      _, params_state = init_fn(None, params)
  
      # Let's modify our params.
      changed_params = tree.map_structure(lambda t: 2.0 * t, params)
      ema_params, params_state = apply_fn_g(
          None, params_state, changed_params
      )
      ema_params2, params_state = apply_fn_g(
          None, params_state, changed_params
      )
  
      # ema_params should be the same as ema_params2 with update_stats=False!
      for p1, p2 in zip(
          tree.flatten(ema_params2), tree.flatten(ema_params)
      ):
          self.assertEqual(p1.shape, p2.shape)
          np.testing.assert_allclose(p1, p2)
  
      def h(x):
          """This will behave like normal."""
          return moving_averages.EMAParamsTree(0.2)(
              x, update_stats=True
          )
  
      init_fn, apply_fn_h = multi_transform.without_apply_rng(
          transform.transform_with_state(h)
      )
      _, params_state = init_fn(None, params)
      params, params_state = apply_fn_h(None, params_state, params)
  
      # Let's modify our params.
      changed_params = tree.map_structure(lambda t: 2.0 * t, params)
      ema_params, params_state = apply_fn_h(
          None, params_state, changed_params
      )
      ema_params2, params_state = apply_fn_h(
          None, params_state, changed_params
      )
  
      # ema_params should be different as ema_params2 with update_stats=False!
      for p1, p2 in zip(
          tree.flatten(ema_params2), tree.flatten(ema_params)
      ):
          self.assertEqual(p1.shape, p2.shape)
          with self.assertRaisesRegex(
              AssertionError, "Not equal to tolerance"
          ):
              np.testing.assert_allclose(p1, p2, atol=1e-6)
  
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 4 for calling the EMAParamsTreeTest obj's underlying member method code with example usage: obj.test_tree_update_stats() and returns: basic, moving_averages, moving_averages