code: def select_direction(
      key: jnp.ndarray, action_prob: jnp.ndarray
  ) -> jnp.int32:
      return random.choice(key, 4, replace=False, p=action_prob.ravel())
  
imports: from typing import Tuple
  from functools import partial
  from PIL import Image
  from PIL import ImageDraw
  import numpy
  import jax
  import jax.numpy
  from jax import random
  from jax import tree_util
  from flax.struct import dataclass
  from evojax.task.base import TaskState
  from evojax.task.base import VectorizedTask
label_desc: Use image node_id 9 for calling a global function with example usage: select_direction(key, action_prob) and returns: random