hunk.code_window:     def test_impala(self):
           check_support("IMPALA", {"num_gpus": 0}, self.stats)
       def test_ppo(self):
           check_support(
               "PPO", {
                   "num_workers": 1,
                   "num_sgd_iter": 1,
                   "train_batch_size": 10,
                   "sample_batch_size": 10,
                   "sgd_minibatch_size": 1,
               },
               self.stats,
               check_bounds=True)
       def test_pg(self):
  
hunk.labels: keep keep keep keep keep replace replace replace replace replace replace replace replace replace replace keep keep
hunk.after_edit:         config = {
               "num_workers": 1,
               "num_sgd_iter": 1,
               "train_batch_size": 10,
               "sample_batch_size": 10,
               "sgd_minibatch_size": 1,
           }
           check_support("PPO", config, self.stats, check_bounds=True)
           config["use_pytorch"] = True
           check_support("PPO", config, self.stats, check_bounds=True)
  
hunk.file_path: rllib/tests/test_supported_spaces.py
hunk.type: replace
file: package org.ray.streaming.api.stream;
  
  import org.ray.streaming.api.function.impl.AggregateFunction;
  import org.ray.streaming.api.function.impl.ReduceFunction;
  import org.ray.streaming.api.partition.impl.KeyPartition;
  import org.ray.streaming.operator.StreamOperator;
  import org.ray.streaming.operator.impl.ReduceOperator;
  
  /**
   * Represents a DataStream returned by a key-by operation.
   *
   * @param <K> Type of the key.
   * @param <T> Type of the data.
   */
  public class KeyDataStream<K, T> extends DataStream<T> {
  
    public KeyDataStream(DataStream<T> input, StreamOperator streamOperator) {
      super(input, streamOperator);
      this.partition = new KeyPartition();
    }
  
    /**
     * Apply a reduce function to this stream.
     *
     * @param reduceFunction The reduce function.
     * @return A new DataStream.
     */
    public DataStream<T> reduce(ReduceFunction reduceFunction) {
      return new DataStream<>(this, new ReduceOperator(reduceFunction));
    }
  
    /**
     * Apply an aggregate Function to this stream.
     *
     * @param aggregateFunction The aggregate function
     * @param <A> The type of aggregated intermediate data.
     * @param <O> The type of result data.
     * @return A new DataStream.
     */
    public <A, O> DataStream<O> aggregate(AggregateFunction<T, A, O> aggregateFunction) {
      return new DataStream<>(this, null);
    }
  
    public KeyDataStream<K, T> setParallelism(int parallelism) {
      this.parallelism = parallelism;
      return this;
    }
  }
  