1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
|
package com.yahoo.searchlib.rankingexpression.integration.onnx.importer.operations;
import com.yahoo.searchlib.rankingexpression.Reference;
import com.yahoo.searchlib.rankingexpression.evaluation.Value;
import com.yahoo.searchlib.rankingexpression.integration.onnx.importer.DimensionRenamer;
import com.yahoo.searchlib.rankingexpression.integration.onnx.importer.OrderedTensorType;
import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
import com.yahoo.searchlib.rankingexpression.rule.ReferenceNode;
import com.yahoo.searchlib.rankingexpression.rule.TensorFunctionNode;
import com.yahoo.tensor.functions.TensorFunction;
import onnx.Onnx;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.function.Function;
/**
* Wraps an ONNX node and produces the respective Vespa tensor operation.
* During import, a graph of these operations are constructed. Then, the
* types are used to deduce sensible dimension names using the
* DimensionRenamer. After the types have been renamed, the proper
* Vespa expressions can be extracted.
*
* @author lesters
*/
public abstract class OnnxOperation {
protected final Onnx.NodeProto node; // can be null for onnx inputs and constants
protected final List<OnnxOperation> inputs;
protected final List<OnnxOperation> outputs = new ArrayList<>();
protected final List<String> importWarnings = new ArrayList<>();
protected OrderedTensorType type;
protected TensorFunction function;
protected Value constantValue = null;
OnnxOperation(Onnx.NodeProto node, List<OnnxOperation> inputs) {
this.node = node;
this.inputs = Collections.unmodifiableList(inputs);
this.inputs.forEach(i -> i.outputs.add(this));
}
protected abstract OrderedTensorType lazyGetType();
protected abstract TensorFunction lazyGetFunction();
/** Returns the Vespa tensor type of this operation if it exists */
public Optional<OrderedTensorType> type() {
if (type == null) {
type = lazyGetType();
}
return Optional.ofNullable(type);
}
/** Returns the Vespa tensor function implementing all operations from this node with inputs */
public Optional<TensorFunction> function() {
if (function == null) {
if (isConstant()) {
ExpressionNode constant = new ReferenceNode(Reference.simple("constant", vespaName()));
function = new TensorFunctionNode.TensorFunctionExpressionNode(constant);
} else {
function = lazyGetFunction();
}
}
return Optional.ofNullable(function);
}
/** Return Onnx node */
public Onnx.NodeProto node() { return node; }
/** Return unmodifiable list of inputs */
public List<OnnxOperation> inputs() { return inputs; }
/** Return unmodifiable list of outputs. If a node has multiple outputs, consider adding a macro. */
public List<OnnxOperation> outputs() { return Collections.unmodifiableList(outputs); }
/** Add dimension name constraints for this operation */
public void addDimensionNameConstraints(DimensionRenamer renamer) { }
/** Performs dimension rename for this operation */
public void renameDimensions(DimensionRenamer renamer) { type = type.rename(renamer); }
/** Return true for operations that are inputs to the model itself (as opposed to inputs to the operation) */
public boolean isInput() { return false; }
/** Return true if this node is constant */
public boolean isConstant() { return inputs.stream().allMatch(OnnxOperation::isConstant); }
/** Gets the constant value if it exists */
public Optional<Value> getConstantValue() { return Optional.ofNullable(constantValue); }
/** Retrieve the valid Vespa name of this node */
public String vespaName() { return vespaName(node.getName()); }
public String vespaName(String name) { return name != null ? namePartOf(name).replace('/', '_') : null; }
/** Retrieve the list of warnings produced during its lifetime */
public List<String> warnings() { return Collections.unmodifiableList(importWarnings); }
/** Set an input warning */
public void warning(String warning) { importWarnings.add(warning); }
boolean verifyInputs(int expected, Function<OnnxOperation, Optional<?>> func) {
if (inputs.size() != expected) {
throw new IllegalArgumentException("Expected " + expected + " inputs " +
"for '" + node.getName() + "', got " + inputs.size());
}
return inputs.stream().map(func).allMatch(Optional::isPresent);
}
boolean allInputTypesPresent(int expected) {
return verifyInputs(expected, OnnxOperation::type);
}
boolean allInputFunctionsPresent(int expected) {
return verifyInputs(expected, OnnxOperation::function);
}
/**
* A method signature input and output has the form name:index.
* This returns the name part without the index.
*/
public static String namePartOf(String name) {
name = name.startsWith("^") ? name.substring(1) : name;
return name.split(":")[0];
}
/**
* This return the output index part. Indexes are used for nodes with
* multiple outputs.
*/
public static int indexPartOf(String name) {
int i = name.indexOf(":");
return i < 0 ? 0 : Integer.parseInt(name.substring(i + 1));
}
}
|
