From 4cc6962e6fb80faf953dca25b3ba0ab4ae39f538 Mon Sep 17 00:00:00 2001
From: Lutz Roeder <lutzroeder@users.noreply.github.com>
Date: Sat, 29 Jun 2024 09:54:05 -0700
Subject: [PATCH] Add TorchScript test file (#1061)

---
 source/python.js             | 14 +++++++-------
 source/pytorch-metadata.json |  1 +
 source/pytorch.js            |  2 +-
 test/models.json             |  8 ++++++++
 4 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/source/python.js b/source/python.js
index 4f50edbca2..67f02c367e 100644
--- a/source/python.js
+++ b/source/python.js
@@ -5088,7 +5088,7 @@ python.Execution = class {
             }
             const dtype = dtypes.get(obj.dtype.str);
             const strides = obj.strides.map((stride) => stride / obj.itemsize);
-            const storage = execution.invoke('torch.storage._TypedStorage', [obj.size, dtype]);
+            const storage = execution.invoke('torch.storage.TypedStorage', [obj.size, dtype]);
             storage._set_cdata(obj.data);
             const tensor = execution.invoke('torch.Tensor', []);
             tensor.__setstate__([storage, 0, obj.shape, strides]);
@@ -5262,7 +5262,7 @@ python.Execution = class {
             const shape = size;
             dtype = torch._prims_common.dtype_or_default(dtype);
             size = shape.reduce((a, b) => a * b, 1);
-            const storage = execution.invoke('torch.storage._TypedStorage', [size, dtype]);
+            const storage = execution.invoke('torch.storage.TypedStorage', [size, dtype]);
             const tensor = execution.invoke('torch.Tensor', []);
             tensor.__setstate__([storage, 0, shape, stride]);
             return tensor;
@@ -6348,13 +6348,13 @@ python.Execution = class {
                 return storage;
             }
         });
-        this.registerType('torch.storage._UntypedStorage', class extends torch.storage._StorageBase {
+        this.registerType('torch.storage.UntypedStorage', class extends torch.storage._StorageBase {
             constructor() {
                 super();
-                throw new python.Error('_UntypedStorage not implemented.');
+                throw new python.Error('UntypedStorage not implemented.');
             }
         });
-        this.registerType('torch.storage._TypedStorage', class {
+        this.registerType('torch.storage.TypedStorage', class {
             constructor(...args) {
                 if (args.length >= 2 && Number.isInteger(args[0]) && args[1] instanceof torch.dtype) {
                     if (args[3] instanceof torch.device) {
@@ -6363,7 +6363,7 @@ python.Execution = class {
                         [this._size, this._dtype] = args;
                     }
                 } else {
-                    throw new python.Error(`Unsupported _TypedStorage arguments '${JSON.stringify(args)}'.`);
+                    throw new python.Error(`Unsupported TypedStorage arguments '${JSON.stringify(args)}'.`);
                 }
             }
             get device() {
@@ -6408,7 +6408,7 @@ python.Execution = class {
                 return storage;
             }
         });
-        this.registerType('torch.storage._LegacyStorage', class extends torch.storage._TypedStorage {
+        this.registerType('torch.storage._LegacyStorage', class extends torch.storage.TypedStorage {
             constructor() {
                 super();
                 throw new python.Error('_LegacyStorage not implemented.');
diff --git a/source/pytorch-metadata.json b/source/pytorch-metadata.json
index 580dc86f08..dcfa4e2344 100755
--- a/source/pytorch-metadata.json
+++ b/source/pytorch-metadata.json
@@ -8064,6 +8064,7 @@
   },
   {
     "name": "aten::lstm_cell",
+    "category": "Layer",
     "inputs": [
       { "name": "input", "type": "Tensor" },
       { "name": "hx", "type": "Tensor[]" },
diff --git a/source/pytorch.js b/source/pytorch.js
index f5536283e4..60d14a98c8 100644
--- a/source/pytorch.js
+++ b/source/pytorch.js
@@ -3216,7 +3216,7 @@ pytorch.jit.FlatBuffersLoader = class {
         const data = this._module.storage_data[index].data;
         const dtype = this._dtypes.get(metadata.scalar_type);
         const size = data.length / dtype.itemsize();
-        const storage = this._cu.execution.invoke('torch.storage._TypedStorage', [size, dtype]);
+        const storage = this._cu.execution.invoke('torch.storage.TypedStorage', [size, dtype]);
         storage._set_cdata(data);
         const tensor = this._cu.execution.invoke('torch.Tensor', []);
         const shape = Array.from(metadata.sizes);
diff --git a/test/models.json b/test/models.json
index 79d28e168f..652f4cecd3 100644
--- a/test/models.json
+++ b/test/models.json
@@ -5978,6 +5978,14 @@
     "format":   "TorchScript v1.6",
     "link":     "https://github.com/lutzroeder/netron/issues/637"
   },
+  {
+    "type":     "pytorch",
+    "target":   "v1_lj_8000.jit",
+    "source":   "https://github.com/user-attachments/files/16041474/v1_lj_8000.jit.zip[v1_lj_8000.jit]",
+    "format":   "TorchScript v1.6",
+    "error":    "Cannot read properties of undefined (reading 'position')",
+    "link":     "https://github.com/lutzroeder/netron/issues/1061"
+  },
   {
     "type":     "pytorch",
     "target":   "v3_1_ru.pt",