code: def test_mean_random(self):
      x = jnp.ones(10)
      bx = jnp.ones((9, 10))
      ffn = RandomFFN()
      wrapped_ffn = networks.FeedForwardNetwork(
          init=functools.partial(ffn.init, x=x), apply=ffn.apply
      )
      mean_ensemble = ensemble.make_ensemble(
          wrapped_ffn, ensemble.apply_mean, num_networks=3
      )
      key = jax.random.PRNGKey(0)
      params = mean_ensemble.init(key)
      single_output = mean_ensemble.apply(params, x)
      self.assertEqual(single_output.shape, (15,))
      batch_output = mean_ensemble.apply(params, bx)
      # Make sure all rows are equal:
      np.testing.assert_allclose(
          jnp.broadcast_to(batch_output[0], batch_output.shape),
          batch_output,
          atol=1e-5,
          rtol=1e-5,
      )
  
      # Check results explicitly:
      all_members = jnp.concatenate(
          [
              jnp.expand_dims(
                  ffn.apply(
                      jax.tree_map(lambda p, i=i: p[i], params), bx
                  ),
                  axis=0,
              )
              for i in range(3)
          ]
      )
      batch_means = jnp.mean(all_members, axis=0)
      np.testing.assert_allclose(
          batch_output, batch_means, atol=1e-5, rtol=1e-5
      )
  
imports: import functools
  from typing import Any
  from acme.agents.jax.mbop import ensemble
  from acme.jax import networks
  from flax import linen
  import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest
label_desc: Use image node_id 12 for calling the EnsembleTest obj's underlying member method code with example usage: obj.test_mean_random() without return types