class_docstr: CMA-ES stochastic optimizer class with ask-and-tell interface, using JAX backend.
  Args:
      param_size:
          A parameter size.
      pop_size:
          A population size (optional).
          If not specified, a proper value will be inferred from param_size.
      mean:
          Initial mean vector of multi-variate gaussian distributions (optional).
          If specified, it must have a dimension of (param_size, ).
          If not specified, a default value of all-zero array will be used.
      init_stdev:
          Initial standard deviation of covariance matrix (optional).
          If not specified, a default value of 0.1 will be used.
      seed:
          A seed number (optional).
      cov:
          A covariance matrix (optional).
      logger:
          A logging.Logger instance (optional).
          If not specified, a new one will be created.
      enable_numeric_check:
          A bool indicating whether to enable numeric check (optional).
          If True, a numeric check for mean and standard deviation is enabled.
          The default value is False.
class_docstr_tok: ["CMA-ES","stochastic","optimizer","class","with","ask-and-tell","interface",",","using","JAX","backend",".","Args",":","param_size",":","A","parameter","size",".","pop_size",":","A","population","size","(","optional",")",".","If","not","specified",",","a","proper","value","will","be","inferred","from","param_size",".","mean",":","Initial","mean","vector","of","multi-variate","gaussian","distributions","(","optional",")",".","If","specified",",","it","must","have","a","dimension","of","(","param_size",",",")",".","If","not","specified",",","a","default","value","of","all-zero","array","will","be","used",".","init_stdev",":","Initial","standard","deviation","of","covariance","matrix","(","optional",")",".","If","not","specified",",","a","default","value","of","0.1","will","be","used",".","seed",":","A","seed","number","(","optional",")",".","cov",":","A","covariance","matrix","(","optional",")",".","logger",":","A","logging.Logger","instance","(","optional",")",".","If","not","specified",",","a","new","one","will","be","created",".","enable_numeric_check",":","A","bool","indicating","whether","to","enable","numeric","check","(","optional",")",".","If","True",",","a","numeric","check","for","mean","and","standard","deviation","is","enabled",".","The","default","value","is","False","."]
code: def __init__(
      self,
      param_size: int,
      pop_size: Optional[int] = None,
      mean: Optional[jnp.ndarray] = None,
      init_stdev: Optional[float] = 0.1,
      seed: Optional[int] = 0,
      cov: Optional[jnp.ndarray] = None,
      logger: logging.Logger = None,
      enable_numeric_check: Optional[bool] = False,
  ):
      if mean is None:
          mean = jnp.zeros(param_size)
      else:
          mean = ensure_jnp(mean)
          assert mean.shape == (param_size,), (
              f"Both params_size and mean are specified."
              f"In this case,  mean (whose shape is {mean.shape}) must have a dimension of (param_size, )"
              f" (i.e. {(param_size, )}), which is not true."
          )
      mean = ensure_jnp(mean)
  
      if enable_numeric_check:
          mean_max = ensure_jnp(
              _MEAN_MAX_X64
              if jax.config.jax_enable_x64
              else _MEAN_MAX_X32
          )
          sigma_max = ensure_jnp(
              _SIGMA_MAX_X64
              if jax.config.jax_enable_x64
              else _SIGMA_MAX_X32
          )
      else:
          mean_max = ensure_jnp(jnp.inf)
          sigma_max = ensure_jnp(jnp.inf)
  
      if enable_numeric_check:
          assert jnp.all(
              jnp.abs(mean) <= mean_max
          ), f"Abs of all elements of mean vector must be less than {mean_max}"
  
      n_dim = len(mean)
      assert n_dim > 1, "The dimension of mean must be larger than 1"
  
      assert (
          init_stdev > 0
      ), "init_stdev must be non-zero positive value"
      sigma = init_stdev
      sigma = ensure_jnp(sigma)
  
      population_size = pop_size
      if population_size is None:
          population_size = 4 + math.floor(
              3 * math.log(n_dim)
          )  # (eq. 48)
      assert (
          population_size > 0
      ), "pop_size must be non-zero positive value."
  
      mu = population_size // 2
  
      # (eq.49)
      weights_prime = jnp.array(
          [
              math.log((population_size + 1) / 2) - math.log(i + 1)
              for i in range(population_size)
          ]
      )
      mu_eff = (jnp.sum(weights_prime[:mu]) ** 2) / jnp.sum(
          weights_prime[:mu] ** 2
      )
      mu_eff_minus = (jnp.sum(weights_prime[mu:]) ** 2) / jnp.sum(
          weights_prime[mu:] ** 2
      )
  
      # learning rate for the rank-one update
      alpha_cov = 2.0
      c1 = alpha_cov / ((n_dim + 1.3) ** 2 + mu_eff)
      assert isinstance(c1, jnp.ndarray)
      # learning rate for the rank-μ update
      cmu = min(
          1 - c1 - 1e-8,  # 1e-8 is for large pop_size.
          alpha_cov
          * (mu_eff - 2 + 1 / mu_eff)
          / ((n_dim + 2) ** 2 + alpha_cov * mu_eff / 2),
      )
      assert (
          c1 + cmu <= 1
      ), "invalid learning rate for the rank-one and/or rank-μ update"
  
      min_alpha = min(
          1 + c1 / cmu,  # eq.50
          1 + (2 * mu_eff_minus) / (mu_eff + 2),  # eq.51
          (1 - c1 - cmu) / (n_dim * cmu),  # eq.52
      )
  
      # (eq.53)
      positive_sum = jnp.sum(weights_prime[weights_prime > 0])
      negative_sum = jnp.sum(jnp.abs(weights_prime[weights_prime < 0]))
      weights = jnp.where(
          weights_prime >= 0,
          1 / positive_sum * weights_prime,
          min_alpha / negative_sum * weights_prime,
      )
      cm = 1  # (eq. 54)
  
      # learning rate for the cumulation for the step-size control (eq.55)
      c_sigma = (mu_eff + 2) / (n_dim + mu_eff + 5)
      d_sigma = (
          1
          + 2 * max(0, math.sqrt((mu_eff - 1) / (n_dim + 1)) - 1)
          + c_sigma
      )
      assert (
          c_sigma < 1
      ), "invalid learning rate for cumulation for the step-size control"
  
      # learning rate for cumulation for the rank-one update (eq.56)
      cc = (4 + mu_eff / n_dim) / (n_dim + 4 + 2 * mu_eff / n_dim)
      assert (
          cc <= 1
      ), "invalid learning rate for cumulation for the rank-one update"
  
      self.hyper_parameters = _HyperParameters(
          n_dim=n_dim,
          pop_size=population_size,
          mu=mu,
      )
  
      self.coefficients = _Coefficients(
          mu_eff=mu_eff,
          cc=cc,
          c1=c1,
          cmu=cmu,
          c_sigma=c_sigma,
          d_sigma=d_sigma,
          cm=cm,
          # E||N(0, I)|| (p.28)
          chi_n=math.sqrt(n_dim)
          * (1.0 - (1.0 / (4.0 * n_dim)) + 1.0 / (21.0 * (n_dim**2))),
          weights=weights,
          sigma_max=sigma_max,
      )
  
      # evolution path (state)
      if cov is None:
          cov = jnp.eye(n_dim)
      else:
          assert cov.shape == (
              n_dim,
              n_dim,
          ), "Invalid shape of covariance matrix"
  
      self.state = _State(
          p_sigma=jnp.zeros(n_dim),
          pc=jnp.zeros(n_dim),
          mean=mean,
          C=cov,
          sigma=sigma,
          D=None,
          B=None,
          g=0,
          key=jax.random.PRNGKey(seed),
      )
  
      # Below are helper vars.
      self._latest_solutions = None
      if logger is None:
          # Change this name accordingly
          self.logger = create_logger(name="CMA")
      else:
          self.logger = logger
  
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 1 to create a new CMA_ES_JAX object from inherited base classes: NEAlgorithm with example: obj = CMA_ES_JAX(param_size, pop_size, mean, init_stdev, seed, cov, logger, enable_numeric_check)