code: def get_task_reset_keys(
      key: jnp.ndarray,
      test: bool,
      pop_size: int,
      n_tests: int,
      n_repeats: int,
      ma_training: bool,
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      key, subkey = random.split(key=key)
      if ma_training:
          reset_keys = random.split(subkey, n_repeats)
      else:
          if test:
              reset_keys = random.split(subkey, n_tests * n_repeats)
          else:
              reset_keys = random.split(subkey, n_repeats)
              reset_keys = jnp.tile(reset_keys, (pop_size, 1))
      return key, reset_keys
  
imports: import logging
  from functools import partial
  from typing import Tuple
  from typing import Union
  import time
  import jax
  import jax.numpy
  from jax import random
  from jax.tree_util import tree_map
  from evojax.obs_norm import ObsNormalizer
  from evojax.task.base import TaskState
  from evojax.task.base import VectorizedTask
  from evojax.policy.base import PolicyState
  from evojax.policy.base import PolicyNetwork
  from evojax.util import create_logger
label_desc: Use image node_id 1 for calling a global function with example usage: get_task_reset_keys(key, test, pop_size, n_tests, n_repeats, ma_training) and returns: key,  reset_keys