add class comment and test expansion

2 files changed, 13 insertions, 0 deletions

diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/EuclideanDistance.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/EuclideanDistance.java
index 4feddf9f808..25399416c29 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/EuclideanDistance.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/EuclideanDistance.java
@@ -8,6 +8,8 @@ import java.util.List;
 import java.util.Objects;
 
 /**
+ * Convenience for euclidean distance between vectors.
+ * euclidean_distance(a, b, mydim) == sqrt(sum(pow(a-b, 2), mydim))
  * @author arnej
  */
 public class EuclideanDistance<NAMETYPE extends Name> extends CompositeTensorFunction<NAMETYPE> {
diff --git a/vespajlib/src/test/java/com/yahoo/tensor/functions/EuclideanDistanceTestCase.java b/vespajlib/src/test/java/com/yahoo/tensor/functions/EuclideanDistanceTestCase.java
index 9d06c313ecc..4fae432b3ca 100644
--- a/vespajlib/src/test/java/com/yahoo/tensor/functions/EuclideanDistanceTestCase.java
+++ b/vespajlib/src/test/java/com/yahoo/tensor/functions/EuclideanDistanceTestCase.java
@@ -3,6 +3,7 @@ package com.yahoo.tensor.functions;
 
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorType;
+import com.yahoo.tensor.evaluation.VariableTensor;
 import org.junit.Test;
 
 import java.util.List;
@@ -40,4 +41,14 @@ public class EuclideanDistanceTestCase {
         assertEquals(expect, result);
     }
 
+    @Test
+    public void testExpansion() {
+        var tType = TensorType.fromSpec("tensor(vecdim[128])");
+        var a = new VariableTensor<>("left", tType);
+        var b = new VariableTensor<>("right", tType);
+        var op = new EuclideanDistance<>(a, b, "vecdim");
+        assertEquals("map(reduce(map(join(left, right, f(a,b)(a - b)), f(a)(a * a)), sum, vecdim), f(a)(sqrt(a)))",
+                     op.toPrimitive().toString());
+    }
+
 }