1 files changed, 49 insertions, 0 deletions
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/importer/DimensionRenamerTest.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/importer/DimensionRenamerTest.java
new file mode 100644
index 00000000000..4bd28a74d6f
--- /dev/null
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/integration/ml/importer/DimensionRenamerTest.java
@@ -0,0 +1,49 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.searchlib.rankingexpression.integration.ml.importer;
+
+import com.yahoo.searchlib.rankingexpression.integration.ml.importer.DimensionRenamer;
+import org.junit.Test;
+
+import static org.junit.Assert.assertTrue;
+
+public class DimensionRenamerTest {
+
+    @Test
+    public void testMnistRenaming() {
+        DimensionRenamer renamer = new DimensionRenamer();
+
+        renamer.addDimension("first_dimension_of_x");
+        renamer.addDimension("second_dimension_of_x");
+        renamer.addDimension("first_dimension_of_w");
+        renamer.addDimension("second_dimension_of_w");
+        renamer.addDimension("first_dimension_of_b");
+
+        // which dimension to join on matmul
+        renamer.addConstraint("second_dimension_of_x", "first_dimension_of_w", DimensionRenamer::equals, null);
+
+        // other dimensions in matmul can't be equal
+        renamer.addConstraint("first_dimension_of_x", "second_dimension_of_w", DimensionRenamer::lesserThan, null);
+
+        // for efficiency, put dimension to join on innermost
+        renamer.addConstraint("first_dimension_of_x", "second_dimension_of_x", DimensionRenamer::lesserThan, null);
+        renamer.addConstraint("first_dimension_of_w", "second_dimension_of_w", DimensionRenamer::greaterThan, null);
+
+        // bias
+        renamer.addConstraint("second_dimension_of_w", "first_dimension_of_b", DimensionRenamer::equals, null);
+
+        renamer.solve();
+
+        String firstDimensionOfXName = renamer.dimensionNameOf("first_dimension_of_x").get();
+        String secondDimensionOfXName = renamer.dimensionNameOf("second_dimension_of_x").get();
+        String firstDimensionOfWName = renamer.dimensionNameOf("first_dimension_of_w").get();
+        String secondDimensionOfWName = renamer.dimensionNameOf("second_dimension_of_w").get();
+        String firstDimensionOfBName = renamer.dimensionNameOf("first_dimension_of_b").get();
+
+        assertTrue(firstDimensionOfXName.compareTo(secondDimensionOfXName) < 0);
+        assertTrue(firstDimensionOfWName.compareTo(secondDimensionOfWName) > 0);
+        assertTrue(secondDimensionOfXName.compareTo(firstDimensionOfWName) == 0);
+        assertTrue(firstDimensionOfXName.compareTo(secondDimensionOfWName) < 0);
+        assertTrue(secondDimensionOfWName.compareTo(firstDimensionOfBName) == 0);
+    }
+
+}