code: def get_state_dataclass(**states):
      if "feet_contact" in states:
          return BDExtractorState(**states)
      else:
          return State(**states)
  
imports: import sys
  from typing import Optional
  from typing import Tuple
  from typing import Union
  import jax
  import jax.numpy
  from flax.struct import dataclass
  from evojax.task.base import VectorizedTask
  from evojax.task.base import TaskState
  from evojax.task.base import BDExtractor
label_desc: Use image node_id 1 for calling a global function with example usage: get_state_dataclass() and returns: BDExtractorState, State