You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@systemds.apache.org by mb...@apache.org on 2021/12/18 20:27:58 UTC
[systemds] branch main updated: [SYSTEMDS-3125] Extended federated tests w/ verified federated output
This is an automated email from the ASF dual-hosted git repository.
mboehm7 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/systemds.git
The following commit(s) were added to refs/heads/main by this push:
new d9dadf1 [SYSTEMDS-3125] Extended federated tests w/ verified federated output
d9dadf1 is described below
commit d9dadf12517597cea7f2af14e3c7b6fe11b82fcc
Author: ywcb00 <yw...@ywcb.org>
AuthorDate: Sat Dec 18 21:22:20 2021 +0100
[SYSTEMDS-3125] Extended federated tests w/ verified federated output
Closes #1419.
---
.../federated/primitives/FederatedCastToMatrixTest.java | 1 +
.../federated/primitives/FederatedCtableTest.java | 8 +++++---
.../federated/primitives/FederatedFullCumulativeTest.java | 3 +++
.../federated/primitives/FederatedRCBindTest.java | 14 +++++++-------
.../primitives/FederatedWeightedDivMatrixMultTest.java | 8 ++++++++
.../federated/primitives/FederatedWeightedSigmoidTest.java | 3 ++-
.../primitives/FederatedWeightedUnaryMatrixMultTest.java | 4 +++-
.../functions/federated/FederatedCtableFedOutput.dml | 5 ++++-
.../federated/FederatedCtableFedOutputReference.dml | 4 +++-
.../functions/federated/cumulative/FederatedCummaxTest.dml | 8 ++++++--
.../federated/cumulative/FederatedCummaxTestReference.dml | 5 +++--
.../functions/federated/cumulative/FederatedCumminTest.dml | 8 ++++++--
.../federated/cumulative/FederatedCumminTestReference.dml | 5 +++--
.../federated/cumulative/FederatedCumprodTest.dml | 8 ++++++--
.../federated/cumulative/FederatedCumprodTestReference.dml | 5 +++--
.../functions/federated/cumulative/FederatedCumsumTest.dml | 7 +++++--
.../federated/cumulative/FederatedCumsumTestReference.dml | 5 +++--
.../federated/cumulative/FederatedCumsumprodTest.dml | 8 ++++++--
.../cumulative/FederatedCumsumprodTestReference.dml | 5 +++--
.../federated/primitives/FederatedCastToMatrixTest.dml | 4 ++++
.../primitives/FederatedCastToMatrixTestReference.dml | 3 +++
.../federated/quaternary/FederatedWDivMMBasicMultTest.dml | 9 +++++++--
.../quaternary/FederatedWDivMMBasicMultTestReference.dml | 2 +-
.../federated/quaternary/FederatedWDivMMLeftMultTest.dml | 9 +++++++--
.../quaternary/FederatedWDivMMLeftMultTestReference.dml | 2 +-
.../federated/quaternary/FederatedWDivMMRightEpsTest.dml | 8 +++++++-
.../quaternary/FederatedWDivMMRightEpsTestReference.dml | 3 ++-
.../quaternary/FederatedWDivMMRightMultMinus4Test.dml | 10 +++++++---
.../FederatedWDivMMRightMultMinus4TestReference.dml | 2 +-
.../quaternary/FederatedWDivMMRightMultMinusTest.dml | 8 +++++++-
.../FederatedWDivMMRightMultMinusTestReference.dml | 3 ++-
.../federated/quaternary/FederatedWDivMMRightMultTest.dml | 8 +++++++-
.../quaternary/FederatedWDivMMRightMultTestReference.dml | 3 ++-
.../federated/quaternary/FederatedWDivMMRightTest.dml | 8 +++++++-
.../quaternary/FederatedWDivMMRightTestReference.dml | 3 ++-
.../federated/quaternary/FederatedWSigmoidLogTest.dml | 9 +++++++--
.../quaternary/FederatedWSigmoidLogTestReference.dml | 2 +-
.../federated/quaternary/FederatedWSigmoidMinusLogTest.dml | 9 +++++++--
.../quaternary/FederatedWSigmoidMinusLogTestReference.dml | 2 +-
.../federated/quaternary/FederatedWSigmoidMinusTest.dml | 9 +++++++--
.../quaternary/FederatedWSigmoidMinusTestReference.dml | 2 +-
.../federated/quaternary/FederatedWSigmoidTest.dml | 9 +++++++--
.../quaternary/FederatedWSigmoidTestReference.dml | 2 +-
.../federated/quaternary/FederatedWUMMExpDivTest.dml | 9 +++++++--
.../quaternary/FederatedWUMMExpDivTestReference.dml | 2 +-
.../federated/quaternary/FederatedWUMMExpMultTest.dml | 9 +++++++--
.../quaternary/FederatedWUMMExpMultTestReference.dml | 2 +-
.../federated/quaternary/FederatedWUMMMult2Test.dml | 9 +++++++--
.../quaternary/FederatedWUMMMult2TestReference.dml | 2 +-
.../federated/quaternary/FederatedWUMMPow2Test.dml | 9 +++++++--
.../quaternary/FederatedWUMMPow2TestReference.dml | 2 +-
51 files changed, 211 insertions(+), 76 deletions(-)
diff --git a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedCastToMatrixTest.java b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedCastToMatrixTest.java
index 4fb95c6..2a0018d 100644
--- a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedCastToMatrixTest.java
+++ b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedCastToMatrixTest.java
@@ -149,6 +149,7 @@ public class FederatedCastToMatrixTest extends AutomatedTestBase {
fedOut = fedOut.split("SystemDS Statistics:")[0];
Assert.assertTrue("Equal Printed Output", out.equals(fedOut));
Assert.assertTrue("Contains federated Cast to frame", heavyHittersContainsString("fed_castdtm"));
+ Assert.assertTrue(heavyHittersContainsString("fed_uak+")); // verify output is federated
TestUtils.shutdownThreads(t1, t2);
rtplatform = platformOld;
diff --git a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedCtableTest.java b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedCtableTest.java
index 9aeb776..ab04efc 100644
--- a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedCtableTest.java
+++ b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedCtableTest.java
@@ -114,7 +114,7 @@ public class FederatedCtableTest extends AutomatedTestBase {
runFedCtable(HOME, TEST_NAME, matrixInput, port1, port2, port3, port4);
else
runNonFedCtable(HOME, TEST_NAME, matrixInput, port1, port2, port3, port4);
- checkResults();
+ checkResults(fedOutput);
TestUtils.shutdownThreads(t1, t2, t3, t4);
resetExecMode(platformOld);
@@ -195,12 +195,14 @@ public class FederatedCtableTest extends AutomatedTestBase {
runTest(true, false, null, -1);
}
- void checkResults() {
+ void checkResults(boolean fedOutput) {
// compare via files
- compareResults(1e-9);
+ compareResults(0);
// check for federated operations
Assert.assertTrue(heavyHittersContainsString("fed_ctable"));
+ if(fedOutput) // verify output is federated
+ Assert.assertTrue(heavyHittersContainsString("fed_uak+"));
// check that federated input files are still existing
Assert.assertTrue(HDFSTool.existsFileOnHDFS(input("X1")));
diff --git a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedFullCumulativeTest.java b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedFullCumulativeTest.java
index aded5fe..72fa3d7 100644
--- a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedFullCumulativeTest.java
+++ b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedFullCumulativeTest.java
@@ -200,6 +200,9 @@ public class FederatedFullCumulativeTest extends AutomatedTestBase {
break;
}
+ if(instType != ExecType.SPARK) // verify output is federated
+ Assert.assertTrue(heavyHittersContainsString("fed_uak+"));
+
// check that federated input files are still existing
Assert.assertTrue(HDFSTool.existsFileOnHDFS(input("X1")));
Assert.assertTrue(HDFSTool.existsFileOnHDFS(input("X2")));
diff --git a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedRCBindTest.java b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedRCBindTest.java
index 04b668d..a3f37f7 100644
--- a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedRCBindTest.java
+++ b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedRCBindTest.java
@@ -23,7 +23,7 @@ import java.util.Arrays;
import java.util.Collection;
import org.apache.sysds.api.DMLScript;
-import org.apache.sysds.common.Types;
+import org.apache.sysds.common.Types.ExecMode;
import org.apache.sysds.runtime.meta.MatrixCharacteristics;
import org.apache.sysds.test.AutomatedTestBase;
import org.apache.sysds.test.TestConfiguration;
@@ -75,17 +75,17 @@ public class FederatedRCBindTest extends AutomatedTestBase {
@Test
public void federatedRCBindCP() {
- federatedRCBind(Types.ExecMode.SINGLE_NODE);
+ federatedRCBind(ExecMode.SINGLE_NODE);
}
@Test
public void federatedRCBindSP() {
- federatedRCBind(Types.ExecMode.SPARK);
+ federatedRCBind(ExecMode.SPARK);
}
- public void federatedRCBind(Types.ExecMode execMode) {
+ public void federatedRCBind(ExecMode execMode) {
boolean sparkConfigOld = DMLScript.USE_LOCAL_SPARK_CONFIG;
- Types.ExecMode platformOld = rtplatform;
+ ExecMode platformOld = rtplatform;
getAndLoadTestConfiguration(TEST_NAME);
String HOME = SCRIPT_DIR + TEST_DIR;
@@ -115,7 +115,7 @@ public class FederatedRCBindTest extends AutomatedTestBase {
Thread t4 = startLocalFedWorkerThread(port4);
// we need the reference file to not be written to hdfs, so we get the correct format
- rtplatform = Types.ExecMode.SINGLE_NODE;
+ rtplatform = ExecMode.SINGLE_NODE;
// Run reference dml script with normal matrix for Row/Col sum
fullDMLScriptName = HOME + TEST_NAME + "Reference.dml";
programArgs = new String[] {"-nvargs", "in_A1=" + input("A1"), "in_A2=" + input("A2"),
@@ -128,7 +128,7 @@ public class FederatedRCBindTest extends AutomatedTestBase {
// reference file should not be written to hdfs, so we set platform here
rtplatform = execMode;
- if(rtplatform == Types.ExecMode.SPARK) {
+ if(rtplatform == ExecMode.SPARK) {
DMLScript.USE_LOCAL_SPARK_CONFIG = true;
}
TestConfiguration config = availableTestConfigurations.get(TEST_NAME);
diff --git a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedDivMatrixMultTest.java b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedDivMatrixMultTest.java
index dd02e3d..f3f2c0f 100644
--- a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedDivMatrixMultTest.java
+++ b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedDivMatrixMultTest.java
@@ -22,6 +22,7 @@ package org.apache.sysds.test.functions.federated.primitives;
import org.apache.sysds.common.Types.ExecMode;
import org.apache.sysds.runtime.matrix.data.MatrixValue.CellIndex;
import org.apache.sysds.runtime.meta.MatrixCharacteristics;
+import org.apache.sysds.runtime.util.CollectionUtils;
import org.apache.sysds.runtime.util.HDFSTool;
import org.apache.sysds.test.AutomatedTestBase;
import org.apache.sysds.test.TestConfiguration;
@@ -36,6 +37,7 @@ import org.junit.runners.Parameterized;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
+import java.util.Set;
@RunWith(value = Parameterized.class)
@net.jcip.annotations.NotThreadSafe
@@ -63,6 +65,10 @@ public class FederatedWeightedDivMatrixMultTest extends AutomatedTestBase
private final static int BLOCKSIZE = 1024;
+ private final static Set<String> FEDERATED_OUTPUT = CollectionUtils.asSet(
+ RIGHT_TEST_NAME, RIGHT_EPS_TEST_NAME, BASIC_MULT_TEST_NAME, LEFT_MULT_TEST_NAME,
+ RIGHT_MULT_TEST_NAME, RIGHT_MULT_MINUS_TEST_NAME, RIGHT_MULT_MINUS_4_TEST_NAME);
+
@Parameterized.Parameter()
public int rows;
@Parameterized.Parameter(1)
@@ -309,6 +315,8 @@ public class FederatedWeightedDivMatrixMultTest extends AutomatedTestBase
// check for federated operations
Assert.assertTrue(heavyHittersContainsString("fed_wdivmm"));
+ if(FEDERATED_OUTPUT.contains(test_name)) // verify the output is federated
+ Assert.assertTrue(heavyHittersContainsString("fed_uak+"));
// check that federated input files are still existing
Assert.assertTrue(HDFSTool.existsFileOnHDFS(input("X1")));
diff --git a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedSigmoidTest.java b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedSigmoidTest.java
index 0ee07bb..d592d20 100644
--- a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedSigmoidTest.java
+++ b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedSigmoidTest.java
@@ -48,7 +48,7 @@ public class FederatedWeightedSigmoidTest extends AutomatedTestBase {
private final static String OUTPUT_NAME = "Z";
- private final static double TOLERANCE = 1e-14;
+ private final static double TOLERANCE = 1e-10;
private final static int BLOCKSIZE = 1024;
@@ -192,6 +192,7 @@ public class FederatedWeightedSigmoidTest extends AutomatedTestBase {
// check for federated operations
Assert.assertTrue(heavyHittersContainsString("fed_wsigmoid", 1, exec_mode == ExecMode.SPARK ? 2 : 3));
+ Assert.assertTrue(heavyHittersContainsString("fed_uak+", 1, 3)); // verify output is federated
// check that federated input files are still existing
Assert.assertTrue(HDFSTool.existsFileOnHDFS(input("X1")));
diff --git a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedUnaryMatrixMultTest.java b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedUnaryMatrixMultTest.java
index 9086edf..acebbf5 100644
--- a/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedUnaryMatrixMultTest.java
+++ b/src/test/java/org/apache/sysds/test/functions/federated/primitives/FederatedWeightedUnaryMatrixMultTest.java
@@ -49,7 +49,7 @@ public class FederatedWeightedUnaryMatrixMultTest extends AutomatedTestBase
private final static String OUTPUT_NAME = "Z";
- private final static double TOLERANCE = 1e-14;
+ private final static double TOLERANCE = 1e-9;
private final static int BLOCKSIZE = 1024;
@@ -188,6 +188,8 @@ public class FederatedWeightedUnaryMatrixMultTest extends AutomatedTestBase
// check for federated operations
Assert.assertTrue(heavyHittersContainsString("fed_wumm", 1, exec_mode == ExecMode.SPARK ? 2 : 3));
+ // verify output is federated
+ Assert.assertTrue(heavyHittersContainsString("fed_uak+", 1, 3));
// check that federated input files are still existing
Assert.assertTrue(HDFSTool.existsFileOnHDFS(input("X1")));
diff --git a/src/test/scripts/functions/federated/FederatedCtableFedOutput.dml b/src/test/scripts/functions/federated/FederatedCtableFedOutput.dml
index a2eda9d..448c919 100644
--- a/src/test/scripts/functions/federated/FederatedCtableFedOutput.dml
+++ b/src/test/scripts/functions/federated/FederatedCtableFedOutput.dml
@@ -50,4 +50,7 @@ else
else
X2 = table(rix, cix);
-write(X2, $out);
+while(FALSE) { }
+Z = as.matrix(sum(X2));
+
+write(Z, $out);
diff --git a/src/test/scripts/functions/federated/FederatedCtableFedOutputReference.dml b/src/test/scripts/functions/federated/FederatedCtableFedOutputReference.dml
index 4fc6852..79ed98e 100644
--- a/src/test/scripts/functions/federated/FederatedCtableFedOutputReference.dml
+++ b/src/test/scripts/functions/federated/FederatedCtableFedOutputReference.dml
@@ -49,4 +49,6 @@ else
else
X2 = table(rix, cix);
-write(X2, $8);
+Z = as.matrix(sum(X2));
+
+write(Z, $8);
diff --git a/src/test/scripts/functions/federated/cumulative/FederatedCummaxTest.dml b/src/test/scripts/functions/federated/cumulative/FederatedCummaxTest.dml
index 9f0f40c..adb0a23 100644
--- a/src/test/scripts/functions/federated/cumulative/FederatedCummaxTest.dml
+++ b/src/test/scripts/functions/federated/cumulative/FederatedCummaxTest.dml
@@ -29,5 +29,9 @@ if ($rP) {
list(0,$cols/2), list($rows, 3*($cols/4)), list(0, 3*($cols/4)), list($rows, $cols)));
}
-s = cummax(A);
-write(s, $out_S);
+S = cummax(A);
+
+while(FALSE) { }
+sS = as.matrix(sum(S));
+
+write(sS, $out_S);
diff --git a/src/test/scripts/functions/federated/cumulative/FederatedCummaxTestReference.dml b/src/test/scripts/functions/federated/cumulative/FederatedCummaxTestReference.dml
index 558c843..43ab960 100644
--- a/src/test/scripts/functions/federated/cumulative/FederatedCummaxTestReference.dml
+++ b/src/test/scripts/functions/federated/cumulative/FederatedCummaxTestReference.dml
@@ -22,5 +22,6 @@
if($6) { A = rbind(read($1), read($2), read($3), read($4)); }
else { A = cbind(read($1), read($2), read($3), read($4)); }
-s = cummax(A);
-write(s, $5);
+S = cummax(A);
+sS = as.matrix(sum(S));
+write(sS, $5);
diff --git a/src/test/scripts/functions/federated/cumulative/FederatedCumminTest.dml b/src/test/scripts/functions/federated/cumulative/FederatedCumminTest.dml
index f088e6b..f983bab 100644
--- a/src/test/scripts/functions/federated/cumulative/FederatedCumminTest.dml
+++ b/src/test/scripts/functions/federated/cumulative/FederatedCumminTest.dml
@@ -29,5 +29,9 @@ if ($rP) {
list(0,$cols/2), list($rows, 3*($cols/4)), list(0, 3*($cols/4)), list($rows, $cols)));
}
-s = cummin(A);
-write(s, $out_S);
+S = cummin(A);
+
+while(FALSE) { }
+sS = as.matrix(sum(S));
+
+write(sS, $out_S);
diff --git a/src/test/scripts/functions/federated/cumulative/FederatedCumminTestReference.dml b/src/test/scripts/functions/federated/cumulative/FederatedCumminTestReference.dml
index e65f584..2025044 100644
--- a/src/test/scripts/functions/federated/cumulative/FederatedCumminTestReference.dml
+++ b/src/test/scripts/functions/federated/cumulative/FederatedCumminTestReference.dml
@@ -22,5 +22,6 @@
if($6) { A = rbind(read($1), read($2), read($3), read($4)); }
else { A = cbind(read($1), read($2), read($3), read($4)); }
-s = cummin(A);
-write(s, $5);
+S = cummin(A);
+sS = as.matrix(sum(S));
+write(sS, $5);
diff --git a/src/test/scripts/functions/federated/cumulative/FederatedCumprodTest.dml b/src/test/scripts/functions/federated/cumulative/FederatedCumprodTest.dml
index 167cbbf..efb3cfa 100644
--- a/src/test/scripts/functions/federated/cumulative/FederatedCumprodTest.dml
+++ b/src/test/scripts/functions/federated/cumulative/FederatedCumprodTest.dml
@@ -29,5 +29,9 @@ if ($rP) {
list(0,$cols/2), list($rows, 3*($cols/4)), list(0, 3*($cols/4)), list($rows, $cols)));
}
-s = cumprod(A);
-write(s, $out_S);
+S = cumprod(A);
+
+while(FALSE) { }
+sS = as.matrix(sum(S));
+
+write(sS, $out_S);
diff --git a/src/test/scripts/functions/federated/cumulative/FederatedCumprodTestReference.dml b/src/test/scripts/functions/federated/cumulative/FederatedCumprodTestReference.dml
index a6c7311..5060a2c 100644
--- a/src/test/scripts/functions/federated/cumulative/FederatedCumprodTestReference.dml
+++ b/src/test/scripts/functions/federated/cumulative/FederatedCumprodTestReference.dml
@@ -22,5 +22,6 @@
if($6) { A = rbind(read($1), read($2), read($3), read($4)); }
else { A = cbind(read($1), read($2), read($3), read($4)); }
-s = cumprod(A);
-write(s, $5);
+S = cumprod(A);
+sS = as.matrix(sum(S));
+write(sS, $5);
diff --git a/src/test/scripts/functions/federated/cumulative/FederatedCumsumTest.dml b/src/test/scripts/functions/federated/cumulative/FederatedCumsumTest.dml
index 757a9d7..3f4d8ac 100644
--- a/src/test/scripts/functions/federated/cumulative/FederatedCumsumTest.dml
+++ b/src/test/scripts/functions/federated/cumulative/FederatedCumsumTest.dml
@@ -29,5 +29,8 @@ if ($rP) {
list(0,$cols/2), list($rows, 3*($cols/4)), list(0, 3*($cols/4)), list($rows, $cols)));
}
-s = cumsum(A);
-write(s, $out_S);
+S = cumsum(A);
+
+while(FALSE) { }
+sS = as.matrix(sum(S));
+write(sS, $out_S);
diff --git a/src/test/scripts/functions/federated/cumulative/FederatedCumsumTestReference.dml b/src/test/scripts/functions/federated/cumulative/FederatedCumsumTestReference.dml
index c496905..99b03c8 100644
--- a/src/test/scripts/functions/federated/cumulative/FederatedCumsumTestReference.dml
+++ b/src/test/scripts/functions/federated/cumulative/FederatedCumsumTestReference.dml
@@ -22,5 +22,6 @@
if($6) { A = rbind(read($1), read($2), read($3), read($4)); }
else { A = cbind(read($1), read($2), read($3), read($4)); }
-s = cumsum(A);
-write(s, $5);
+S = cumsum(A);
+sS = as.matrix(sum(S));
+write(sS, $5);
diff --git a/src/test/scripts/functions/federated/cumulative/FederatedCumsumprodTest.dml b/src/test/scripts/functions/federated/cumulative/FederatedCumsumprodTest.dml
index 43978d9..a646505 100644
--- a/src/test/scripts/functions/federated/cumulative/FederatedCumsumprodTest.dml
+++ b/src/test/scripts/functions/federated/cumulative/FederatedCumsumprodTest.dml
@@ -28,5 +28,9 @@ if ($rP) {
ranges=list(list(0, 0), list($rows, $cols/4), list(0,$cols/4), list($rows, $cols/2),
list(0,$cols/2), list($rows, 3*($cols/4)), list(0, 3*($cols/4)), list($rows, $cols)));
}
-s = cumsumprod(A[,1:2]);
-write(s, $out_S);
+S = cumsumprod(A[,1:2]);
+
+while(FALSE) { }
+sS = as.matrix(sum(S));
+
+write(sS, $out_S);
diff --git a/src/test/scripts/functions/federated/cumulative/FederatedCumsumprodTestReference.dml b/src/test/scripts/functions/federated/cumulative/FederatedCumsumprodTestReference.dml
index 5c73b14..0f7610e 100644
--- a/src/test/scripts/functions/federated/cumulative/FederatedCumsumprodTestReference.dml
+++ b/src/test/scripts/functions/federated/cumulative/FederatedCumsumprodTestReference.dml
@@ -21,5 +21,6 @@
if($6) { A = rbind(read($1), read($2), read($3), read($4)); }
else { A = cbind(read($1), read($2), read($3), read($4)); }
-s = cumsumprod(A[,1:2]);
-write(s, $5);
+S = cumsumprod(A[,1:2]);
+sS = as.matrix(sum(S));
+write(sS, $5);
diff --git a/src/test/scripts/functions/federated/primitives/FederatedCastToMatrixTest.dml b/src/test/scripts/functions/federated/primitives/FederatedCastToMatrixTest.dml
index 52b9889..c82db78 100644
--- a/src/test/scripts/functions/federated/primitives/FederatedCastToMatrixTest.dml
+++ b/src/test/scripts/functions/federated/primitives/FederatedCastToMatrixTest.dml
@@ -24,3 +24,7 @@ X = federated(type="frame", addresses=list($X1, $X2),
Z = as.matrix(X)
print(toString(Z[1]))
+
+while(FALSE) { }
+sZ = sum(Z);
+print("Sum: " + sZ);
diff --git a/src/test/scripts/functions/federated/primitives/FederatedCastToMatrixTestReference.dml b/src/test/scripts/functions/federated/primitives/FederatedCastToMatrixTestReference.dml
index a0db27d..791fe2b 100644
--- a/src/test/scripts/functions/federated/primitives/FederatedCastToMatrixTestReference.dml
+++ b/src/test/scripts/functions/federated/primitives/FederatedCastToMatrixTestReference.dml
@@ -23,3 +23,6 @@ X = rbind(read($X1), read($X2))
Z = as.matrix(X)
print(toString(Z[1]))
+
+sZ = sum(Z);
+print("Sum: " + sZ);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMBasicMultTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMBasicMultTest.dml
index 72e5616..6b608a2 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMBasicMultTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMBasicMultTest.dml
@@ -36,8 +36,13 @@ X = t(X); # col partitioned X
while(FALSE) { }
Z3 = X * (V %*% t(U));
-while(FALSE) { }
-Z = (Z1 + Z2) + sum(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z)
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMBasicMultTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMBasicMultTestReference.dml
index 6b5d1cb..8bd63be 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMBasicMultTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMBasicMultTestReference.dml
@@ -33,6 +33,6 @@ X = t(X);
Z3 = X * (V %*% t(U));
-Z = (Z1 + Z2) + sum(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z)
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMLeftMultTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMLeftMultTest.dml
index 03b9f90..b856f0d 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMLeftMultTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMLeftMultTest.dml
@@ -36,8 +36,13 @@ X = t(X); # col partitioned X
while(FALSE) { }
Z3 = t(t(V) %*% (X * (V %*% t(U))));
-while(FALSE) { }
-Z = Z1 + Z2 + sum(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMLeftMultTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMLeftMultTestReference.dml
index 03eb263..8fe9e0e 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMLeftMultTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMLeftMultTestReference.dml
@@ -33,6 +33,6 @@ X = t(X);
Z3 = t(t(V) %*% (X * (V %*% t(U))));
-Z = Z1 + Z2 + sum(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightEpsTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightEpsTest.dml
index a457ff1..b92fc94 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightEpsTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightEpsTest.dml
@@ -26,6 +26,12 @@ U = read($in_U)
V = read($in_V)
eps = $in_W
-Z = (X / (U %*% t(V) + eps)) %*% V;
+tmp = (X / (U %*% t(V) + eps)) %*% V;
+
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ = sum(tmp);
+while(FALSE) { }
+Z = as.matrix(sZ);
write(Z, $out_Z)
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightEpsTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightEpsTestReference.dml
index 1c9efa2..0cbd6ac 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightEpsTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightEpsTestReference.dml
@@ -24,6 +24,7 @@ U = read($in_U)
V = read($in_V)
eps = $in_W
-Z = (X / (U %*% t(V) + eps)) %*% V;
+tmp = (X / (U %*% t(V) + eps)) %*% V;
+Z = as.matrix(sum(tmp));
write(Z, $out_Z)
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinus4Test.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinus4Test.dml
index 687e607..50864d3 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinus4Test.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinus4Test.dml
@@ -39,9 +39,13 @@ MX = t(MX); # col partitioned federated MX
while(FALSE) { }
Z3 = (X * (V %*% t(U) - MX)) %*% U;
-while(FALSE) { }
-
-Z = Z1 + Z2 + sum(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinus4TestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinus4TestReference.dml
index 00a0a5a..092632a 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinus4TestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinus4TestReference.dml
@@ -36,6 +36,6 @@ MX = t(MX);
Z3 = (X * (V %*% t(U) - MX)) %*% U;
-Z = Z1 + Z2 + sum(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinusTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinusTest.dml
index 203415e..029d618 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinusTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinusTest.dml
@@ -25,6 +25,12 @@ X = federated(addresses=list($in_X1, $in_X2),
U = read($in_U)
V = read($in_V)
-Z = ((X != 0) * (U %*% t(V) - X)) %*% V;
+tmp = ((X != 0) * (U %*% t(V) - X)) %*% V;
+
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ = sum(tmp);
+while(FALSE) { }
+Z = as.matrix(sZ);
write(Z, $out_Z)
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinusTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinusTestReference.dml
index 8f114d1..e3c5f25 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinusTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultMinusTestReference.dml
@@ -23,6 +23,7 @@ X = rbind(read($in_X1), read($in_X2))
U = read($in_U)
V = read($in_V)
-Z = ((X != 0) * (U %*% t(V) - X)) %*% V;
+tmp = ((X != 0) * (U %*% t(V) - X)) %*% V;
+Z = as.matrix(sum(tmp));
write(Z, $out_Z)
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultTest.dml
index cab65d5..d758af3 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultTest.dml
@@ -25,6 +25,12 @@ X = federated(addresses=list($in_X1, $in_X2),
U = read($in_U)
V = read($in_V)
-Z = (X * (U %*% t(V))) %*% V;
+tmp = (X * (U %*% t(V))) %*% V;
+
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ = sum(tmp);
+while(FALSE) { }
+Z = as.matrix(sZ);
write(Z, $out_Z)
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultTestReference.dml
index 535b7f5..06ef31b 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightMultTestReference.dml
@@ -23,6 +23,7 @@ X = rbind(read($in_X1), read($in_X2))
U = read($in_U)
V = read($in_V)
-Z = (X * (U %*% t(V))) %*% V;
+tmp = (X * (U %*% t(V))) %*% V;
+Z = as.matrix(sum(tmp));
write(Z, $out_Z)
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightTest.dml
index de64cf1..f53b8e8 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightTest.dml
@@ -25,6 +25,12 @@ X = federated(addresses=list($in_X1, $in_X2),
U = read($in_U)
V = read($in_V)
-Z = (X / (U %*% t(V))) %*% V;
+tmp = (X / (U %*% t(V))) %*% V;
+
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ = sum(tmp);
+while(FALSE) { }
+Z = as.matrix(sZ);
write(Z, $out_Z)
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightTestReference.dml
index 1fa64e8..6355910 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWDivMMRightTestReference.dml
@@ -23,6 +23,7 @@ X = rbind(read($in_X1), read($in_X2))
U = read($in_U)
V = read($in_V)
-Z = (X / (U %*% t(V))) %*% V;
+tmp = (X / (U %*% t(V))) %*% V;
+Z = as.matrix(sum(tmp));
write(Z, $out_Z)
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidLogTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidLogTest.dml
index 2008d7e..ee875b0 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidLogTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidLogTest.dml
@@ -39,8 +39,13 @@ while(FALSE) { }
UV = V %*% t(U);
Z3 = X * log(1 / (1 + exp(-UV)));
-while(FALSE) { }
-Z = Z1 + Z2 + mean(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidLogTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidLogTestReference.dml
index cf3e28d..cbe3d4b 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidLogTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidLogTestReference.dml
@@ -36,6 +36,6 @@ X = t(X);
UV = V %*% t(U);
Z3 = X * log(1 / (1 + exp(-UV)));
-Z = Z1 + Z2 + mean(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusLogTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusLogTest.dml
index cd806b2..ced6b98 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusLogTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusLogTest.dml
@@ -39,8 +39,13 @@ while(FALSE) { }
UV = -(V %*% t(U));
Z3 = X * log(1 / (1 + exp(-UV)));
-while(FALSE) { }
-Z = Z1 + Z2 + mean(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusLogTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusLogTestReference.dml
index c04c71b..19b462f 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusLogTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusLogTestReference.dml
@@ -36,6 +36,6 @@ X = t(X);
UV = -(V %*% t(U));
Z3 = X * log(1 / (1 + exp(-UV)));
-Z = Z1 + Z2 + mean(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusTest.dml
index d1d0cab..4685f00 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusTest.dml
@@ -39,8 +39,13 @@ while(FALSE) { }
UV = -(V %*% t(U));
Z3 = X * (1 / (1 + exp(-UV)));
-while(FALSE) { }
-Z = Z1 + Z2 + mean(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusTestReference.dml
index 5385c78..bc561e2 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidMinusTestReference.dml
@@ -36,6 +36,6 @@ X = t(X);
UV = -(V %*% t(U));
Z3 = X * (1 / (1 + exp(-UV)));
-Z = Z1 + Z2 + mean(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidTest.dml
index 3162eaa..a70cff3 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidTest.dml
@@ -39,8 +39,13 @@ while(FALSE) { }
UV = V %*% t(U);
Z3 = X * (1 / (1 + exp(-UV)));
-while(FALSE) { }
-Z = Z1 + Z2 + mean(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidTestReference.dml
index 7dff33f..c2b0f4d 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWSigmoidTestReference.dml
@@ -36,6 +36,6 @@ X = t(X);
UV = V %*% t(U);
Z3 = X * (1 / (1 + exp(-UV)));
-Z = Z1 + Z2 + mean(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpDivTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpDivTest.dml
index 2ac851a..b8d3d45 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpDivTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpDivTest.dml
@@ -36,8 +36,13 @@ X = t(X); # col partitioned X
while(FALSE) { }
Z3 = X / exp(V %*% t(U));
-while(FALSE) { }
-Z = Z1 + Z2 + mean(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpDivTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpDivTestReference.dml
index 7083c22..e260784 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpDivTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpDivTestReference.dml
@@ -33,6 +33,6 @@ X = t(X);
Z3 = X / exp(V %*% t(U));
-Z = Z1 + Z2 + mean(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpMultTest.dml b/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpMultTest.dml
index 03fe08a..1ade7d7 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpMultTest.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpMultTest.dml
@@ -36,8 +36,13 @@ X = t(X); # col partitioned X
while(FALSE) { }
Z3 = X * exp(V %*% t(U));
-while(FALSE) { }
-Z = Z1 + Z2 + mean(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpMultTestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpMultTestReference.dml
index c01a244..77b38a3 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpMultTestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWUMMExpMultTestReference.dml
@@ -33,6 +33,6 @@ X = t(X);
Z3 = X * exp(V %*% t(U));
-Z = Z1 + Z2 + mean(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWUMMMult2Test.dml b/src/test/scripts/functions/federated/quaternary/FederatedWUMMMult2Test.dml
index 8126e79..d31b34e 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWUMMMult2Test.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWUMMMult2Test.dml
@@ -36,8 +36,13 @@ X = t(X); # col paritioned X
while(FALSE) { }
Z3 = X * (2 * (V %*% t(U)));
-while(FALSE) { }
-Z = Z1 + Z2 + mean(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWUMMMult2TestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWUMMMult2TestReference.dml
index a99e0d7..7838a13 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWUMMMult2TestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWUMMMult2TestReference.dml
@@ -33,6 +33,6 @@ X = t(X);
Z3 = X * (2 * (V %*% t(U)));
-Z = Z1 + Z2 + mean(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWUMMPow2Test.dml b/src/test/scripts/functions/federated/quaternary/FederatedWUMMPow2Test.dml
index 8c9642f..c0a08c0 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWUMMPow2Test.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWUMMPow2Test.dml
@@ -36,8 +36,13 @@ X = t(X); # col partitioned X
while(FALSE) { }
Z3 = X / (V %*% t(U))^2;
-while(FALSE) { }
-Z = Z1 + Z2 + mean(Z3);
+while(FALSE) { }
+# verify output is federated (creating fed_uak+ heavy hitter)
+sZ1 = sum(Z1);
+sZ2 = sum(Z2);
+sZ3 = sum(Z3);
+while(FALSE) { }
+Z = as.matrix(sZ1 - sZ2 + sZ3);
write(Z, $out_Z);
diff --git a/src/test/scripts/functions/federated/quaternary/FederatedWUMMPow2TestReference.dml b/src/test/scripts/functions/federated/quaternary/FederatedWUMMPow2TestReference.dml
index 6e454e7..66751ff 100644
--- a/src/test/scripts/functions/federated/quaternary/FederatedWUMMPow2TestReference.dml
+++ b/src/test/scripts/functions/federated/quaternary/FederatedWUMMPow2TestReference.dml
@@ -33,6 +33,6 @@ X = t(X);
Z3 = X / (V %*% t(U))^2;
-Z = Z1 + Z2 + mean(Z3);
+Z = as.matrix(sum(Z1) - sum(Z2) + sum(Z3));
write(Z, $out_Z);