Fix method meta error in wasm build (pytorch#13496)

### Summary
An error would occur when trying to get the metadata for a method
containing an input or output that isn't a tensor. This occurs because
there wasn't a check for the input/output tag when generating the list
of tensor metadata.

Now the list of tensor metadata will have `undefined` if the
input/output in that index is not a tensor.

### Test plan
Added unit tests to test the method metadata for a method that has ints
in its input and outputs.
Added unit test to check if getMethod works if the module has multiple
methods.

Unit tests are in the CI but can be ran with
```
bash scripts/build_wasm_tests.sh
cd cmake-out-wasm/extension/wasm/test/
npm test # after installing Jest
```

Author: Conarnar

extension/wasm/test/CMakeLists.txt

@@ -11,6 +11,13 @@
 
 set(MODELS_DIR ${CMAKE_CURRENT_BINARY_DIR}/models/)
 
+add_custom_command(
+  OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/models/test.pte
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/../../..
+  COMMAND python3 -m extension.wasm.test.test_model
+          ${CMAKE_CURRENT_BINARY_DIR}/models/test.pte
+)
+
 add_custom_command(
   OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/models/add_mul.pte
          ${CMAKE_CURRENT_BINARY_DIR}/models/add.pte
@@ -23,8 +30,9 @@ add_custom_command(
 )
 
 add_custom_target(
-  executorch_wasm_test_models DEPENDS ${MODELS_DIR}/add_mul.pte
-                                      ${MODELS_DIR}/add.pte
+  executorch_wasm_test_models
+  DEPENDS ${MODELS_DIR}/add_mul.pte ${MODELS_DIR}/add.pte
+          ${MODELS_DIR}/test.pte
 )
 
 add_custom_command(

extension/wasm/test/test_model.py

@@ -0,0 +1,34 @@
+import sys
+
+import torch
+from executorch.exir import to_edge_transform_and_lower
+from torch.export import export
+
+
+class IndexModel(torch.nn.Module):
+    def forward(self, x, n):
+        return x[n]
+
+
+class AddAllModel(torch.nn.Module):
+    def forward(self, x, n):
+        return x, n, x + n
+
+
+if __name__ == "__main__":
+    output_filepath = sys.argv[1] if len(sys.argv) > 1 else "test.pte"
+    indexModel = IndexModel().eval()
+    addAllModel = AddAllModel().eval()
+
+    exported_index = export(indexModel, (torch.randn([3]), 1))
+    exported_add_all = export(addAllModel, (torch.randn([2, 2]), 1))
+    edge = to_edge_transform_and_lower(
+        {
+            "forward": exported_index,
+            "index": exported_index,
+            "add_all": exported_add_all,
+        }
+    )
+    et = edge.to_executorch()
+    with open(output_filepath, "wb") as file:
+        file.write(et.buffer)

extension/wasm/test/unittests.js

@@ -105,6 +105,13 @@ describe("Module", () => {
         module.delete();
     });
 
+    test("multiple methods", () => {
+        const module = et.Module.load("test.pte");
+        const methods = module.getMethods();
+        expect(methods).toEqual(expect.arrayContaining(["forward", "index", "add_all"]));
+        module.delete();
+    });
+
     test("loadMethod forward", () => {
         const module = et.Module.load("add.pte");
         expect(() => module.loadMethod("forward")).not.toThrow();
@@ -224,6 +231,25 @@ describe("Module", () => {
                 });
                 module.delete();
             });
+
+            test("non-tensor in input", () => {
+                const module = et.Module.load("test.pte");
+                const methodMeta = module.getMethodMeta("add_all");
+                expect(methodMeta.inputTags).toEqual([et.Tag.Tensor, et.Tag.Int]);
+                expect(methodMeta.inputTensorMeta[0]).not.toBeUndefined();
+                expect(methodMeta.inputTensorMeta[1]).toBeUndefined();
+                module.delete();
+            });
+
+            test("non-tensor in output", () => {
+                const module = et.Module.load("test.pte");
+                const methodMeta = module.getMethodMeta("add_all");
+                expect(methodMeta.outputTags).toEqual([et.Tag.Tensor, et.Tag.Int, et.Tag.Tensor]);
+                expect(methodMeta.outputTensorMeta[0]).not.toBeUndefined();
+                expect(methodMeta.outputTensorMeta[1]).toBeUndefined();
+                expect(methodMeta.outputTensorMeta[2]).not.toBeUndefined();
+                module.delete();
+            });
         });
     });
 

extension/wasm/wasm_bindings.cpp

@@ -475,16 +475,26 @@ struct ET_EXPERIMENTAL JsMethodMeta {
         val::array(),
         meta.num_instructions()};
     for (int i = 0; i < meta.num_inputs(); i++) {
-      js_array_push(new_meta.input_tags, meta.input_tag(i).get());
-      js_array_push(
-          new_meta.input_tensor_meta,
-          JsTensorInfo::from_tensor_info(meta.input_tensor_meta(i).get()));
+      Tag tag = meta.input_tag(i).get();
+      js_array_push(new_meta.input_tags, tag);
+      if (tag == Tag::Tensor) {
+        js_array_push(
+            new_meta.input_tensor_meta,
+            JsTensorInfo::from_tensor_info(meta.input_tensor_meta(i).get()));
+      } else {
+        js_array_push(new_meta.input_tensor_meta, val::undefined());
+      }
     }
     for (int i = 0; i < meta.num_outputs(); i++) {
-      js_array_push(new_meta.output_tags, meta.output_tag(i).get());
-      js_array_push(
-          new_meta.output_tensor_meta,
-          JsTensorInfo::from_tensor_info(meta.output_tensor_meta(i).get()));
+      Tag tag = meta.output_tag(i).get();
+      js_array_push(new_meta.output_tags, tag);
+      if (tag == Tag::Tensor) {
+        js_array_push(
+            new_meta.output_tensor_meta,
+            JsTensorInfo::from_tensor_info(meta.output_tensor_meta(i).get()));
+      } else {
+        js_array_push(new_meta.output_tensor_meta, val::undefined());
+      }
     }
     for (int i = 0; i < meta.num_attributes(); i++) {
       js_array_push(