code: def test_apply_round_robin(self):
      x = jnp.ones(10)  # Base input
      bx = jnp.ones((7, 10))  # Batched input
  
      wrapped_ffn = params_adding_ffn(x)
  
      rr_ensemble = ensemble.make_ensemble(
          wrapped_ffn, ensemble.apply_round_robin, num_networks=3
      )
      key = jax.random.PRNGKey(0)
      params = rr_ensemble.init(key)
      self.assertTupleEqual(params.shape, (3,) + x.shape)
  
      y = rr_ensemble.apply(params, jnp.broadcast_to(x, (3,) + x.shape))
      self.assertTupleEqual(y.shape, (3,) + x.shape)
      np.testing.assert_allclose(params, y - x)
  
      # Note: the ensemble dimension must lead, the batch dimension is no longer
      # the leading dimension.
      by = rr_ensemble.apply(
          params,
          jnp.broadcast_to(
              jnp.expand_dims(bx, axis=0), (3,) + bx.shape
          ),
      )
      self.assertTupleEqual(by.shape, (3,) + bx.shape)
  
      # If num_networks=3, then `round_robin(params, input)[4]` should be equal
      # to `apply(params[1], input[4])`, etc.
      yy = rr_ensemble.apply(
          params, jnp.broadcast_to(x, (6,) + x.shape)
      )
      self.assertTupleEqual(yy.shape, (6,) + x.shape)
      np.testing.assert_allclose(
          jnp.concatenate([params, params], axis=0),
          yy - jnp.expand_dims(x, axis=0),
      )
  
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 3 for calling the EnsembleTest obj's underlying member method code with example usage: obj.test_apply_round_robin() without return types