code: def _tell_core(
      hps: _HyperParameters,
      coeff: _Coefficients,
      state: _State,
      fitness: jnp.ndarray,
      solutions: jnp.ndarray,
  ) -> _State:
      next_state = state
  
      g = state.g + 1
      next_state = next_state._replace(g=g)
  
      ranking = jnp.argsort(fitness, axis=0)
  
      sorted_solutions = solutions[ranking]
  
      B, D = state.B, state.D
  
      # Sample new population of search_points, for k=1, ..., pop_size
      B, D = state.B, state.D  # already computed.
      next_state = next_state._replace(B=None, D=None)
  
      x_k = jnp.array(sorted_solutions)  # ~ N(m, σ^2 C)
      y_k = (x_k - state.mean) / state.sigma  # ~ N(0, C)
  
      # Selection and recombination
      # y_w = jnp.sum(y_k[: hps.mu].T * coeff.weights[: hps.mu], axis=1)  # eq.41
      # use lax.dynamic_slice_in_dim here:
      y_w = jnp.sum(
          jax.lax.dynamic_slice_in_dim(y_k, 0, hps.mu, axis=0).T
          * jax.lax.dynamic_slice_in_dim(
              coeff.weights, 0, hps.mu, axis=0
          ),
          axis=1,
      )
      mean = state.mean + coeff.cm * state.sigma * y_w
      next_state = next_state._replace(mean=mean)
  
      # Step-size control
      C_2 = B.dot(jnp.diag(1 / D)).dot(B.T)  # C^(-1/2) = B D^(-1) B^T
      p_sigma = (1 - coeff.c_sigma) * state.p_sigma + jnp.sqrt(
          coeff.c_sigma * (2 - coeff.c_sigma) * coeff.mu_eff
      ) * C_2.dot(y_w)
      next_state = next_state._replace(p_sigma=p_sigma)
  
      norm_p_sigma = jnp.linalg.norm(state.p_sigma)
      sigma = state.sigma * jnp.exp(
          (coeff.c_sigma / coeff.d_sigma)
          * (norm_p_sigma / coeff.chi_n - 1)
      )
      sigma = jnp.min(jnp.array([sigma, coeff.sigma_max]))
      next_state = next_state._replace(sigma=sigma)
  
      # Covariance matrix adaption
      h_sigma_cond_left = norm_p_sigma / jnp.sqrt(
          1 - (1 - coeff.c_sigma) ** (2 * (state.g + 1))
      )
      h_sigma_cond_right = (1.4 + 2 / (hps.n_dim + 1)) * coeff.chi_n
      # h_sigma = 1.0 if h_sigma_cond_left < h_sigma_cond_right else 0.0  # (p.28)
      h_sigma = jax.lax.cond(
          pred=(h_sigma_cond_left < h_sigma_cond_right),
          true_fun=(lambda: 1.0),
          false_fun=(lambda: 0.0),
      )
  
      # (eq.45)
      pc = (1 - coeff.cc) * state.pc + h_sigma * jnp.sqrt(
          coeff.cc * (2 - coeff.cc) * coeff.mu_eff
      ) * y_w
      next_state = next_state._replace(pc=pc)
  
      # (eq.46)
      w_io = coeff.weights * jnp.where(
          coeff.weights >= 0,
          1,
          hps.n_dim
          / (jnp.linalg.norm(C_2.dot(y_k.T), axis=0) ** 2 + _EPS),
      )
  
      delta_h_sigma = (
          (1 - h_sigma) * coeff.cc * (2 - coeff.cc)
      )  # (p.28)
      # assert delta_h_sigma <= 1
  
      # (eq.47)
      rank_one = jnp.outer(state.pc, state.pc)
  
      # This way of computing rank_mu can lead to OOM:
      # rank_mu = jnp.sum(
      #    jnp.array([w * jnp.outer(y, y) for w, y in zip(w_io, y_k)]), axis=0
      # )
      # Try another way of computing rank_mu
      rank_mu = jnp.zeros_like(state.C)
      for w, y in zip(w_io, y_k):
          rank_mu = rank_mu + w * jnp.outer(y, y)
  
      C = (
          (
              1
              + coeff.c1 * delta_h_sigma
              - coeff.c1
              - coeff.cmu * jnp.sum(coeff.weights)
          )
          * state.C
          + coeff.c1 * rank_one
          + coeff.cmu * rank_mu
      )
      next_state = next_state._replace(C=C)
  
      return next_state
  
imports: from typing import NamedTuple
  from functools import partial
  import logging
  import math
  from typing import Optional
  import numpy
  import jax
  from jax import numpy
  from evojax.algo.base import NEAlgorithm
  from evojax.util import create_logger
label_desc: Use image node_id 2 for calling a global function with example usage: _tell_core(hps, coeff, state, fitness, solutions) and returns: next_state