class_docstr: A convolutional neural network for the MNIST classification task.
class_docstr_tok: ["A","convolutional","neural","network","for","the","MNIST","classification","task","."]
code: def __init__(self, logger: logging.Logger = None):
      if logger is None:
          self._logger = create_logger("ConvNetPolicy")
      else:
          self._logger = logger
  
      model = CNN()
      params = model.init(random.PRNGKey(0), jnp.zeros([1, 28, 28, 1]))
      self.num_params, format_params_fn = get_params_format_fn(params)
      self._logger.info(
          "ConvNetPolicy.num_params = {}".format(self.num_params)
      )
      self._format_params_fn = jax.vmap(format_params_fn)
      self._forward_fn = jax.vmap(model.apply)
  
file_path: evojax/evojax/policy/convnet.py
imports: import logging
  from typing import Tuple
  import jax
  import jax.numpy
  from jax import random
  from flax import linen
  from evojax.policy.base import PolicyNetwork
  from evojax.policy.base import PolicyState
  from evojax.task.base import TaskState
  from evojax.util import create_logger
  from evojax.util import get_params_format_fn
label_desc: Use image node_id 1 to create a new ConvNetPolicy object from inherited base classes: PolicyNetwork with example: obj = ConvNetPolicy(logger)