@Test
public void testICQLinearRegression() throws Exception {

	final int m = 10000;
	final int n = 3;
	List<Tuple2<DenseVector, Double>> data = new ArrayList<>();
	for (int i = 0; i < m; ++i) {
		DenseVector feature = DenseVector.rand(n);
		data.add(Tuple2.of(feature.append(1.0), feature.dot(DenseVector.ones(n))));
	}
	DataSet<Tuple2<DenseVector, Double>> trainData =
		MLEnvironmentFactory
			.getDefault()
			.getExecutionEnvironment()
			.fromCollection(data);
	DataSet<DenseVector> initialCoefs =
		MLEnvironmentFactory
			.getDefault()
			.getExecutionEnvironment()
			.fromCollection(Collections.singletonList(DenseVector.zeros(n + 1)));
	DataSet<Double> sampleCount = DataSetUtils
		.countElementsPerPartition(trainData)
		.sum(1)
		.map(new MapFunction<Tuple2<Integer, Long>, Double>() {
			@Override
			public Double map(Tuple2<Integer, Long> value) throws Exception {
				return value.f1.doubleValue();
			}
		});
	final double learningRate = 1.0;
	DataSet<Row> model = new IterativeComQueue()
		.setMaxIter(100)
		.initWithPartitionedData(TRAIN_DATA, trainData)
		.initWithBroadcastData(COEFS, initialCoefs)
		.initWithBroadcastData(SAMPLE_COUNT, sampleCount)
		.add(new ComputeFunction() {
			@Override
			public void calc(ComContext context) {
				List<Tuple2<DenseVector, Double>> trainData = context.getObj(TRAIN_DATA);
				List<DenseVector> coefs = context.getObj(COEFS);
				double[] grads = context.getObj("grads");
				if (grads == null) {
					grads = new double[coefs.get(0).size()];
					context.putObj("grads", grads);
				}
				Arrays.fill(grads, 0.0);
				DenseVector gradsWrapper = new DenseVector(grads);
				for (Tuple2<DenseVector, Double> sample : trainData) {
					gradsWrapper.plusScaleEqual(sample.f0, sample.f1 - sample.f0.dot(coefs.get(0)));
				}
			}
		})
		.add(new AllReduce("grads"))
		.add(new ComputeFunction() {
			@Override
			public void calc(ComContext context) {
				List<DenseVector> coefs = context.getObj(COEFS);
				double[] grads = context.getObj("grads");
				List<Double> sampleCount = context.getObj(SAMPLE_COUNT);
				coefs.get(0).plusScaleEqual(new DenseVector(grads), learningRate / sampleCount.get(0));
			}
		})
		.closeWith(new CompleteResultFunction() {
			@Override
			public List<Row> calc(ComContext context) {
				if (context.getTaskId() == 0) {
					List<DenseVector> coefs = context.getObj(COEFS);
					return Collections.singletonList(Row.of(coefs.get(0)));
				} else {
					return null;
				}
			}
		})
		.exec();
	List<Row> modelL = model.collect();
	Assert.assertEquals(data.get(0).f1,
		data.get(0).f0.dot((Vector) modelL.get(0).getField(0)),
		2.0
	);
}