Add non-primitive tensor expand function

3 files changed, 97 insertions, 1 deletions

diff --git a/vespajlib/abi-spec.json b/vespajlib/abi-spec.json
index e68a37b15b6..8387a3d1a83 100644
--- a/vespajlib/abi-spec.json
+++ b/vespajlib/abi-spec.json
@@ -1699,6 +1699,21 @@
     ],
     "fields": []
   },
+  "com.yahoo.tensor.functions.Expand": {
+    "superClass": "com.yahoo.tensor.functions.CompositeTensorFunction",
+    "interfaces": [],
+    "attributes": [
+      "public"
+    ],
+    "methods": [
+      "public void <init>(com.yahoo.tensor.functions.TensorFunction, java.lang.String)",
+      "public java.util.List arguments()",
+      "public com.yahoo.tensor.functions.TensorFunction withArguments(java.util.List)",
+      "public com.yahoo.tensor.functions.PrimitiveTensorFunction toPrimitive()",
+      "public java.lang.String toString(com.yahoo.tensor.functions.ToStringContext)"
+    ],
+    "fields": []
+  },
   "com.yahoo.tensor.functions.Generate": {
     "superClass": "com.yahoo.tensor.functions.PrimitiveTensorFunction",
     "interfaces": [],
@@ -2053,6 +2068,22 @@
     ],
     "fields": []
   },
+  "com.yahoo.tensor.functions.ScalarFunctions$Constant": {
+    "superClass": "java.lang.Object",
+    "interfaces": [
+      "java.util.function.Function"
+    ],
+    "attributes": [
+      "public"
+    ],
+    "methods": [
+      "public void <init>(double)",
+      "public java.lang.Double apply(java.util.List)",
+      "public java.lang.String toString()",
+      "public bridge synthetic java.lang.Object apply(java.lang.Object)"
+    ],
+    "fields": []
+  },
   "com.yahoo.tensor.functions.ScalarFunctions$Cos": {
     "superClass": "java.lang.Object",
     "interfaces": [
@@ -2603,7 +2634,8 @@
       "public static java.util.function.DoubleUnaryOperator selu(double, double)",
       "public static java.util.function.Function random()",
       "public static java.util.function.Function equal(java.util.List)",
-      "public static java.util.function.Function sum(java.util.List)"
+      "public static java.util.function.Function sum(java.util.List)",
+      "public static java.util.function.Function constant(double)"
     ],
     "fields": []
   },
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Expand.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Expand.java
new file mode 100644
index 00000000000..8fc246a7d9d
--- /dev/null
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Expand.java
@@ -0,0 +1,48 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.tensor.functions;
+
+import com.yahoo.tensor.TensorType;
+import com.yahoo.tensor.evaluation.Name;
+
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * The <i>expand</i> tensor function returns a tensor with a new dimension of
+ * size 1 is added, equivalent to "tensor * tensor(dim_name[1])(1)".
+ *
+ * @author lesters
+ */
+public class Expand<NAMETYPE extends Name> extends CompositeTensorFunction<NAMETYPE> {
+
+    private final TensorFunction<NAMETYPE> argument;
+    private final String dimensionName;
+
+    public Expand(TensorFunction<NAMETYPE> argument, String dimension) {
+        this.argument = argument;
+        this.dimensionName = dimension;
+    }
+
+    @Override
+    public List<TensorFunction<NAMETYPE>> arguments() { return Collections.singletonList(argument); }
+
+    @Override
+    public TensorFunction<NAMETYPE> withArguments(List<TensorFunction<NAMETYPE>> arguments) {
+        if (arguments.size() != 1)
+            throw new IllegalArgumentException("Expand must have 1 argument, got " + arguments.size());
+        return new Expand<>(arguments.get(0), dimensionName);
+    }
+
+    @Override
+    public PrimitiveTensorFunction<NAMETYPE> toPrimitive() {
+        TensorType type = new TensorType.Builder(TensorType.Value.INT8).indexed(dimensionName, 1).build();
+        Generate<NAMETYPE> expansion = new Generate<>(type, ScalarFunctions.constant(1.0));
+        return new Join<>(expansion, argument, ScalarFunctions.multiply());
+    }
+
+    @Override
+    public String toString(ToStringContext context) {
+        return "expand(" + argument.toString(context) + ", " + dimensionName + ")";
+    }
+
+}
diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/ScalarFunctions.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/ScalarFunctions.java
index d6fcd17b8fb..4c2b64244e5 100644
--- a/vespajlib/src/main/java/com/yahoo/tensor/functions/ScalarFunctions.java
+++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/ScalarFunctions.java
@@ -66,6 +66,7 @@ public class ScalarFunctions {
     public static Function<List<Long>, Double> random() { return new Random(); }
     public static Function<List<Long>, Double> equal(List<String> argumentNames) { return new EqualElements(argumentNames); }
     public static Function<List<Long>, Double> sum(List<String> argumentNames) { return new SumElements(argumentNames); }
+    public static Function<List<Long>, Double> constant(double value) { return new Constant(value); }
 
     // Binary operators -----------------------------------------------------------------------------
 
@@ -493,4 +494,19 @@ public class ScalarFunctions {
         }
     }
 
+    public static class Constant implements Function<List<Long>, Double> {
+        private final double value;
+
+        public Constant(double value) {
+            this.value = value;
+        }
+        @Override
+        public Double apply(List<Long> values) {
+            return value;
+        }
+        @Override
+        public String toString() { return Double.toString(value); }
+    }
+
+
 }