code: def test_cma_es_jax_save_and_load_state(self):
      from evojax.algo import CMA_ES_JAX
      from jax import numpy as jnp
  
      solver = CMA_ES_JAX(pop_size=16, param_size=16)
      # one step
      _ = solver.ask()
      solver.tell(jnp.arange(16, dtype=jnp.float32))
      # one step
      state = solver.save_state()
      internal_state_0 = solver.state._asdict()
      # one step
      _ = solver.ask()
      solver.tell(-jnp.arange(16, dtype=jnp.float32))
      internal_state_1 = solver.state._asdict()
      solver.load_state(state)
      internal_state_2 = solver.state._asdict()
  
      keys = list(internal_state_0.keys())
  
      assert not all(
          [
              jnp.all(internal_state_0[key] == internal_state_1[key])
              for key in keys
          ]
      )
      assert all(
          [
              jnp.all(internal_state_0[key] == internal_state_2[key])
              for key in keys
          ]
      )
  
label_desc: Use image node_id 1 for calling the TestAlgo obj's underlying member method code with example usage: obj.test_cma_es_jax_save_and_load_state() without return types