code: def update_score_and_mask(score, reward, mask, done):
      new_score = score + reward * mask
      new_mask = mask * (1 - done.ravel())
      return new_score, new_mask
  
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 7 for calling a global function with example usage: update_score_and_mask(score, reward, mask, done) and returns: new_score,  new_mask