Find sidebar type/shape support (#542)

lutzroeder · Nov 28, 2020 · 2695c8e · 2695c8e
1 parent 5520298
commit 2695c8e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/source/onnx.js b/source/onnx.js
@@ -1104,7 +1104,7 @@ onnx.Utility = class {
             map[onnx.proto.TensorProto.DataType.INT32] = 'int32';
             map[onnx.proto.TensorProto.DataType.INT64] = 'int64';
             map[onnx.proto.TensorProto.DataType.STRING] = 'string';
-            map[onnx.proto.TensorProto.DataType.BOOL] = 'bool';
+            map[onnx.proto.TensorProto.DataType.BOOL] = 'boolean';
             map[onnx.proto.TensorProto.DataType.FLOAT16] = 'float16';
             map[onnx.proto.TensorProto.DataType.DOUBLE] = 'float64';
             map[onnx.proto.TensorProto.DataType.UINT32] = 'uint32';

diff --git a/source/view-sidebar.js b/source/view-sidebar.js
@@ -1057,9 +1057,9 @@ sidebar.DocumentationSidebar = class {
 
 sidebar.FindSidebar = class {
 
-    constructor(host, graphElement, graph) {
+    constructor(host, element, graph) {
         this._host = host;
-        this._graphElement = graphElement;
+        this._graphElement = element;
         this._graph = graph;
         this._contentElement = this._host.document.createElement('div');
         this._contentElement.setAttribute('class', 'sidebar-view-find');
@@ -1155,7 +1155,12 @@ sidebar.FindSidebar = class {
 
             for (const input of node.inputs) {
                 for (const argument of input.arguments) {
-                    if (argument.name && argument.name.toLowerCase().indexOf(text) != -1 && !edgeMatches.has(argument.name)) {
+                    const match =
+                        (argument.name && argument.name.toLowerCase().indexOf(text) != -1) ||
+                        (argument.type && argument.type.dataType && argument.type.dataType.toLowerCase() === text) ||
+                        (argument.type && argument.type.shape && argument.type.shape.dimensions && argument.type.shape.dimensions.some((dimension) => dimension.toString() === text)) ||
+                        (argument.type && argument.type.shape && argument.type.shape.dimensions && (argument.type.shape.toString() === text || argument.type.shape.toString() === '[' + text + ']'));
+                    if (match && !edgeMatches.has(argument.name)) {
                         if (!argument.initializer) {
                             const inputItem = this._host.document.createElement('li');
                             inputItem.innerText = '\u2192 ' + argument.name.split('\n').shift(); // custom argument id