Fix TensorFlow dimension renaming at join

2 files changed, 5 insertions, 3 deletions

diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/OrderedTensorType.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/OrderedTensorType.java
index 3742e443a06..db762d5ddb0 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/OrderedTensorType.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/OrderedTensorType.java
@@ -42,6 +42,8 @@ public class OrderedTensorType {
         return this.type;
     }
 
+    public int rank() { return dimensions.size(); }
+
     public List<TensorType.Dimension> dimensions() {
         return dimensions;
     }
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/Join.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/Join.java
index aa27ba2684d..3e6e036636d 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/Join.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/integration/tensorflow/importer/operations/Join.java
@@ -61,10 +61,10 @@ public class Join extends TensorFlowOperation {
         // "end" (highest numbered) dimensions of the first, but I'm not sure whether this is generally true.
         // Anyway, we move the dimensions of b to the last dimensions of a (instead of by default, the first).
 
-        TensorType a = inputs.get(0).type().get().type();
-        TensorType b = inputs.get(1).type().get().type();
+        OrderedTensorType a = inputs.get(0).type().get();
+        OrderedTensorType b = inputs.get(1).type().get();
         if (a.rank() < b.rank()) {
-            TensorType temp = a;
+            OrderedTensorType temp = a;
             a = b;
             b = temp;
         }