code: def feature_control_rewards(
      features: Array,
      cumulant_type="absolute_change",
      discount=None,
  ) -> Array:
      """Calculates cumulants for feature control tasks from a sequence of features.
  
      For each feature dimension, a distinct pseudo reward is computed based on the
      change in the feature value between consecutive timesteps. Depending on
      `cumulant_type`, cumulants may be equal the features themselves, the absolute
      difference between their values in consecutive steps, their increase/decrease,
      or may take the form of a potential-based reward discounted by `discount`.
  
      See "Reinforcement Learning with Unsupervised Auxiliary Tasks" by Jaderberg,
      Mnih, Czarnecki et al. (https://arxiv.org/abs/1611.05397).
  
      Args:
        features: A tensor of shape `[T+1,D]` of features.
        cumulant_type: either 'feature' (feature is the reward), `absolute_change`
          (the reward equals the absolute difference between consecutive
          timesteps), `increase` (the reward equals the increase in the
          value of the feature), `decrease` (the reward equals the decrease in the
          value of the feature), or 'potential' (r=gamma*phi_{t+1} - phi_t).
        discount: (optional) discount for potential based rewards.
  
      Returns:
        A tensor of cumulants calculated from the features. The shape is `[T,D]`.
      """
      chex.assert_rank(features, 2)
      chex.assert_type(features, float)
  
      if cumulant_type == "feature":
          return features[1:]
      elif cumulant_type == "absolute_change":
          return jnp.abs(features[1:] - features[:-1])
      elif cumulant_type == "increase":
          return features[1:] - features[:-1]
      elif cumulant_type == "decrease":
          return features[:-1] - features[1:]
      elif cumulant_type == "potential":
          return discount * features[1:] - features[:-1]
      else:
          raise ValueError(f"Unknown cumulant_type {cumulant_type}")
  
label_desc: Use image node_id 2 for calling a global function with example usage: feature_control_rewards(features, cumulant_type, discount) and returns: features, jnp, unknown, unknown, unknown