Update safetensors.js (#1113)

lutzroeder · Jul 4, 2023 · 874fdeb · 874fdeb
1 parent c1b90e3
commit 874fdeb
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/source/safetensors.js b/source/safetensors.js
@@ -100,8 +100,18 @@ safetensors.TensorType = class {
 
     constructor(dtype, shape) {
         switch (dtype) {
-            case 'F16': this.dataType = 'float16'; break;
-            case 'F32': this.dataType = 'float32'; break;
+            case 'I8':   this.dataType = 'int8'; break;
+            case 'I16':  this.dataType = 'int16'; break;
+            case 'I32':  this.dataType = 'int32'; break;
+            case 'I64':  this.dataType = 'int64'; break;
+            case 'U8':   this.dataType = 'uint8'; break;
+            case 'U16':  this.dataType = 'uint16'; break;
+            case 'U32':  this.dataType = 'uint32'; break;
+            case 'U64':  this.dataType = 'uint64'; break;
+            case 'BF16': this.dataType = 'bfloat16'; break;
+            case 'F16':  this.dataType = 'float16'; break;
+            case 'F32':  this.dataType = 'float32'; break;
+            case 'F64':  this.dataType = 'float64'; break;
             default: throw new safetensors.Error("Unsupported data type '" + dtype + "'.");
         }
         this.shape = shape;

diff --git a/test/models.json b/test/models.json
@@ -5418,7 +5418,7 @@
   {
     "type":     "safetensors",
     "target":   "simple.safetensors",
-    "source":   "https://github.com/lutzroeder/netron/files/11858547/simple.safetensors.zip[simple.safetensors]",
+    "source":   "https://github.com/lutzroeder/netron/files/11943528/simple.safetensors.zip[simple.safetensors]",
     "format":   "Safetensors",
     "link":     "https://github.com/lutzroeder/netron/issues/1113"
   },