code: def test_round_robin_random(self):
      x = jnp.ones(10)  # Base input
      bx = jnp.ones((9, 10))  # Batched input
      ffn = RandomFFN()
      wrapped_ffn = networks.FeedForwardNetwork(
          init=functools.partial(ffn.init, x=x), apply=ffn.apply
      )
      rr_ensemble = ensemble.make_ensemble(
          wrapped_ffn, ensemble.apply_round_robin, num_networks=3
      )
  
      key = jax.random.PRNGKey(0)
      params = rr_ensemble.init(key)
      out = rr_ensemble.apply(params, bx)
      # The output should be the same every 3 rows:
      blocks = jnp.split(out, 3, axis=0)
      np.testing.assert_array_equal(blocks[0], blocks[1])
      np.testing.assert_array_equal(blocks[0], blocks[2])
      self.assertTrue((out[0] != out[1]).any())
  
      for i in range(9):
          np.testing.assert_allclose(
              out[i],
              ffn.apply(
                  jax.tree_map(lambda p, i=i: p[i % 3], params), bx[i]
              ),
              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 11 for calling the EnsembleTest obj's underlying member method code with example usage: obj.test_round_robin_random() without return types