diff options
Diffstat (limited to 'vespajlib')
4 files changed, 27 insertions, 15 deletions
diff --git a/vespajlib/abi-spec.json b/vespajlib/abi-spec.json index 66eb4b1f4e6..b8b6716d879 100644 --- a/vespajlib/abi-spec.json +++ b/vespajlib/abi-spec.json @@ -1279,6 +1279,7 @@ "public void <init>()", "public void <init>(com.yahoo.tensor.TensorType$Value)", "public varargs void <init>(com.yahoo.tensor.TensorType[])", + "public varargs void <init>(boolean, com.yahoo.tensor.TensorType[])", "public void <init>(java.lang.Iterable)", "public void <init>(com.yahoo.tensor.TensorType$Value, java.lang.Iterable)", "public int rank()", diff --git a/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java b/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java index aeed8c33093..d8959147ee0 100644 --- a/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java +++ b/vespajlib/src/main/java/com/yahoo/tensor/TensorType.java @@ -323,7 +323,7 @@ public class TensorType { * [N] + [] = [] * [] + {} = {} */ - Dimension combineWith(Optional<Dimension> other) { + Dimension combineWith(Optional<Dimension> other, boolean allowDifferentSizes) { if ( ! other.isPresent()) return this; if (this instanceof MappedDimension) return this; if (other.get() instanceof MappedDimension) return other.get(); @@ -333,7 +333,11 @@ public class TensorType { // both are indexed bound IndexedBoundDimension thisIb = (IndexedBoundDimension)this; IndexedBoundDimension otherIb = (IndexedBoundDimension)other.get(); - return thisIb.size().get() < otherIb.size().get() ? thisIb : otherIb; + if (allowDifferentSizes) + return thisIb.size().get() < otherIb.size().get() ? thisIb : otherIb; + if ( ! thisIb.size().equals(otherIb.size())) + throw new IllegalArgumentException("Unequal dimension sizes in " + thisIb + " and " + otherIb); + return thisIb; } @Override @@ -494,9 +498,13 @@ public class TensorType { * The value type will be the largest of the value types of the input types */ public Builder(TensorType ... types) { + this(true, types); + } + + public Builder(boolean allowDifferentSizes, TensorType ... types) { this.valueType = TensorType.combinedValueType(types); for (TensorType type : types) - addDimensionsOf(type); + addDimensionsOf(type, allowDifferentSizes); } /** Creates a builder from the given dimensions, having double as the value type */ @@ -514,17 +522,17 @@ public class TensorType { private static final boolean supportsMixedTypes = false; - private void addDimensionsOf(TensorType type) { + private void addDimensionsOf(TensorType type, boolean allowDifferentSizes) { if ( ! supportsMixedTypes) { // TODO: Support it - addDimensionsOfAndDisallowMixedDimensions(type); + addDimensionsOfAndDisallowMixedDimensions(type, allowDifferentSizes); } else { for (Dimension dimension : type.dimensions) - set(dimension.combineWith(Optional.ofNullable(dimensions.get(dimension.name())))); + set(dimension.combineWith(Optional.ofNullable(dimensions.get(dimension.name())), allowDifferentSizes)); } } - private void addDimensionsOfAndDisallowMixedDimensions(TensorType type) { + private void addDimensionsOfAndDisallowMixedDimensions(TensorType type, boolean allowDifferentSizes) { boolean containsMapped = dimensions.values().stream().anyMatch(d -> ! d.isIndexed()); containsMapped = containsMapped || type.dimensions().stream().anyMatch(d -> ! d.isIndexed()); @@ -532,7 +540,7 @@ public class TensorType { if (containsMapped) dimension = new MappedDimension(dimension.name()); Dimension existing = dimensions.get(dimension.name()); - set(dimension.combineWith(Optional.ofNullable(existing))); + set(dimension.combineWith(Optional.ofNullable(existing), allowDifferentSizes)); } } diff --git a/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java b/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java index 1e0eaa7fad3..5419d04a4fb 100644 --- a/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java +++ b/vespajlib/src/main/java/com/yahoo/tensor/functions/Join.java @@ -48,7 +48,12 @@ public class Join<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETYP /** Returns the type resulting from applying Join to the two given types */ public static TensorType outputType(TensorType a, TensorType b) { - return new TensorType.Builder(a, b).build(); + try { + return new TensorType.Builder(false, a, b).build(); + } + catch (IllegalArgumentException e) { + throw new IllegalArgumentException("Can not join " + a + " and " + b, e); + } } public DoubleBinaryOperator combinator() { return combinator; } @@ -75,14 +80,14 @@ public class Join<NAMETYPE extends Name> extends PrimitiveTensorFunction<NAMETYP @Override public TensorType type(TypeContext<NAMETYPE> context) { - return new TensorType.Builder(argumentA.type(context), argumentB.type(context)).build(); + return outputType(argumentA.type(context), argumentB.type(context)); } @Override public Tensor evaluate(EvaluationContext<NAMETYPE> context) { Tensor a = argumentA.evaluate(context); Tensor b = argumentB.evaluate(context); - TensorType joinedType = new TensorType.Builder(a.type(), b.type()).build(); + TensorType joinedType = outputType(a.type(), b.type()); return evaluate(a, b, joinedType, combinator); } diff --git a/vespajlib/src/test/java/com/yahoo/tensor/TensorTestCase.java b/vespajlib/src/test/java/com/yahoo/tensor/TensorTestCase.java index b1851b5f120..5bd1bbdba37 100644 --- a/vespajlib/src/test/java/com/yahoo/tensor/TensorTestCase.java +++ b/vespajlib/src/test/java/com/yahoo/tensor/TensorTestCase.java @@ -130,13 +130,11 @@ public class TensorTestCase { assertEquals("Mapped vector", 42, (int)dotProduct(vector(Type.mapped), vectors(Type.mapped, 2))); assertEquals("Indexed unbound vector", 42, (int)dotProduct(vector(3, Type.indexedUnbound), vectors(5, Type.indexedUnbound, 2))); assertEquals("Indexed unbound vector", 42, (int)dotProduct(vector(5, Type.indexedUnbound), vectors(3, Type.indexedUnbound, 2))); - assertEquals("Indexed bound vector", 42, (int)dotProduct(vector(3, Type.indexedBound), vectors(5, Type.indexedBound, 2))); - assertEquals("Indexed bound vector", 42, (int)dotProduct(vector(5, Type.indexedBound), vectors(3, Type.indexedBound, 2))); + assertEquals("Indexed bound vector", 42, (int)dotProduct(vector(3, Type.indexedBound), vectors(3, Type.indexedBound, 2))); assertEquals("Mapped matrix", 42, (int)dotProduct(vector(Type.mapped), matrix(Type.mapped, 2))); assertEquals("Indexed unbound matrix", 42, (int)dotProduct(vector(3, Type.indexedUnbound), matrix(5, Type.indexedUnbound, 2))); assertEquals("Indexed unbound matrix", 42, (int)dotProduct(vector(5, Type.indexedUnbound), matrix(3, Type.indexedUnbound, 2))); - assertEquals("Indexed bound matrix", 42, (int)dotProduct(vector(3, Type.indexedBound), matrix(5, Type.indexedBound, 2))); - assertEquals("Indexed bound matrix", 42, (int)dotProduct(vector(5, Type.indexedBound), matrix(3, Type.indexedBound, 2))); + assertEquals("Indexed bound matrix", 42, (int)dotProduct(vector(3, Type.indexedBound), matrix(3, Type.indexedBound, 2))); assertEquals("Mixed vector", 42, (int)dotProduct(vector(Type.mapped), vectors(Type.indexedUnbound, 2))); assertEquals("Mixed vector", 42, (int)dotProduct(vector(Type.mapped), vectors(Type.indexedUnbound, 2))); assertEquals("Mixed matrix", 42, (int)dotProduct(vector(Type.mapped), matrix(Type.indexedUnbound, 2))); |