Add ONNX test file (onnx/onnx#6318)

lutzroeder · Aug 24, 2024 · cd1752e · cd1752e
1 parent 6b168fc
commit cd1752e
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 5 deletions.
diff --git a/source/base.js b/source/base.js
@@ -106,6 +106,13 @@ if (!DataView.prototype.getBfloat16) {
     DataView.__bfloat16_get_uint16_be = new Uint16Array(DataView.__bfloat16_get_float32_be.buffer, DataView.__bfloat16_get_float32_be.byteOffset, 2);
 }
 
+DataView.__float4e2m1_float32 = new Float32Array([0, 0.5, 1, 1.5, 2, 3, 4, 6, -0, -0.5, -1, -1.5, -2, -3, -4, -6]);
+DataView.prototype.getFloat4e2m1 = function(byteOffset) {
+    let value = this.getUint8(byteOffset >> 1);
+    value = byteOffset & 1 ? value >> 4 : value & 0x0F;
+    return DataView.__float4e2m1_float32[value];
+};
+
 DataView.__float8e4m3_float32 = new Float32Array(1);
 DataView.__float8e4m3_uint32 = new Uint32Array(DataView.__float8e4m3_float32.buffer, DataView.__float8e4m3_float32.byteOffset, 1);
 DataView.prototype.getFloat8e4m3 = function(byteOffset, fn, uz) {
@@ -752,6 +759,10 @@ base.Tensor = class {
                             context.dataType = 'int';
                             context.bits = parseInt(dataType.substring(3), 10);
                             context.itemsize = 1;
+                        } else if (dataType === 'float4e2m1') {
+                            context.dataType = 'float4e2m1';
+                            context.bits = 4;
+                            context.itemsize = 1;
                         } else {
                             throw new Error(`Tensor data type '${dataType}' is not implemented.`);
                         }
@@ -909,6 +920,11 @@ base.Tensor = class {
                         results.push(view.getComplex128(offset, this._littleEndian));
                     }
                     break;
+                case 'float4e2m1':
+                    for (; offset < max; offset += stride) {
+                        results.push(view.getFloat4e2m1(offset));
+                    }
+                    break;
                 case 'float8e4m3fn':
                     for (; offset < max; offset += stride) {
                         results.push(view.getFloat8e4m3(offset, true, false));

diff --git a/source/onnx-proto.js b/source/onnx-proto.js
@@ -12,7 +12,8 @@ onnx.Version = {
     "IR_VERSION_2020_5_8": 7,
     "IR_VERSION_2021_7_30": 8,
     "IR_VERSION_2023_5_5": 9,
-    "IR_VERSION": 10
+    "IR_VERSION_2024_3_25": 10,
+    "IR_VERSION": 11
 };
 
 onnx.AttributeProto = class AttributeProto {
@@ -893,7 +894,8 @@ onnx.TensorProto.DataType = {
     "FLOAT8E5M2": 19,
     "FLOAT8E5M2FNUZ": 20,
     "UINT4": 21,
-    "INT4": 22
+    "INT4": 22,
+    "FLOAT4E2M1": 23
 };
 
 onnx.TensorProto.Segment = class Segment {

diff --git a/source/onnx.js b/source/onnx.js
@@ -571,6 +571,7 @@ onnx.Tensor = class {
                                 this._encoding = '<';
                             }
                             break;
+                        case onnx.DataType.FLOAT4E2M1:
                         case onnx.DataType.FLOAT8E4M3FN:
                         case onnx.DataType.FLOAT8E4M3FNUZ:
                         case onnx.DataType.FLOAT8E5M2:
@@ -1019,7 +1020,8 @@ onnx.DataType = {
     FLOAT8E5M2: 19,
     FLOAT8E5M2FNUZ: 20,
     UINT4: 21,
-    INT4: 22
+    INT4: 22,
+    FLOAT4E2M1: 23
 };
 
 onnx.AttributeType = {

diff --git a/test/models.json b/test/models.json
@@ -3978,8 +3978,16 @@
   },
   {
     "type":     "onnx",
-    "target":   "fp8.onnx",
-    "source":   "https://github.com/lutzroeder/netron/files/11702419/fp8.onnx.zip[fp8.onnx]",
+    "target":   "float4e2m1.pb",
+    "source":   "https://github.com/user-attachments/files/16737906/float4e2m1.pb.zip[float4e2m1.pb]",
+    "format":   "ONNX Tensor",
+    "assert":   "model.graphs[0].nodes[0].attributes[0].value.type.dataType == 'float4e2m1'",
+    "link":     "https://github.com/lutzroeder/netron/issues/6"
+  },
+  {
+    "type":     "onnx",
+    "target":   "float8.onnx",
+    "source":   "https://github.com/user-attachments/files/16737904/float8.onnx.zip[float8.onnx]",
     "format":   "ONNX v9",
     "link":     "https://github.com/lutzroeder/netron/issues/6"
   },