Support negative dimesions in import of expanddims

1 files changed, 7 insertions, 3 deletions

diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/OperationMapper.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/OperationMapper.java
index f77699e8f9e..f9cec4a14f0 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/OperationMapper.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/OperationMapper.java
@@ -281,11 +281,15 @@ class OperationMapper {
         if (axis.type().rank() != 0) {
             throw new IllegalArgumentException("Axis argument to ExpandDims must be a scalar");
         }
-        int dimensionToInsert = (int)axis.asDouble();
 
         TensorFunction inputFunction = arguments.get(0).function();
         TensorType inputType = arguments.get(0).type();
 
+        int dimensionToInsert = (int)axis.asDouble();
+        if (dimensionToInsert < 0) {
+            dimensionToInsert = inputType.dimensions().size() - dimensionToInsert;
+        }
+
         TensorType.Builder outputTypeBuilder = new TensorType.Builder();
         int dimensionIndex = 0;
         for (int i = 0; i < inputType.dimensions().size() + 1; ++i) {
@@ -294,8 +298,8 @@ class OperationMapper {
             if (i == dimensionToInsert) {
                 size = 1L;
             } else {
-                TensorType.Dimension dimension = inputType.dimensions().get(dimensionIndex);
-                size = dimension.size().orElseThrow(() -> new IllegalArgumentException("Dimension has no size"));
+                size = dimensionSize(inputType.dimensions().get(dimensionIndex));
+                dimensionIndex++;
             }
             outputTypeBuilder.indexed(name, size);
         }