Set enum names for convenience

Author: Conarnar

extension/wasm/test/unittests.js

@@ -59,17 +59,15 @@ describe("Tensor", () => {
 
     test("scalar type", () => {
         const tensor = et.Tensor.ones([2, 2]);
-        // ScalarType can only be checked by strict equality.
-        expect(tensor.scalarType).toBe(et.ScalarType.Float);
+        expect(tensor.scalarType).toEqual(et.ScalarType.Float);
         tensor.delete();
     });
 
     test("long tensor", () => {
         const tensor = et.Tensor.ones([2, 2], et.ScalarType.Long);
         expect(tensor.data).toEqual(new BigInt64Array([1n, 1n, 1n, 1n]));
         expect(tensor.sizes).toEqual([2, 2]);
-        // ScalarType can only be checked by strict equality.
-        expect(tensor.scalarType).toBe(et.ScalarType.Long);
+        expect(tensor.scalarType).toEqual(et.ScalarType.Long);
         tensor.delete();
     });
 
@@ -78,8 +76,7 @@ describe("Tensor", () => {
         const tensor = et.Tensor.fromArray([2, 2], [1n, 2n, 3n, 4n]);
         expect(tensor.data).toEqual(new BigInt64Array([1n, 2n, 3n, 4n]));
         expect(tensor.sizes).toEqual([2, 2]);
-        // ScalarType can only be checked by strict equality.
-        expect(tensor.scalarType).toBe(et.ScalarType.Long);
+        expect(tensor.scalarType).toEqual(et.ScalarType.Long);
         tensor.delete();
     });
 });
@@ -124,17 +121,15 @@ describe("Module", () => {
             const module = et.Module.load("add_mul.pte");
             const methodMeta = module.getMethodMeta("forward");
             expect(methodMeta.inputTags.length).toEqual(3);
-            // Tags can only be checked by strict equality.
-            methodMeta.inputTags.forEach((tag) => expect(tag).toBe(et.Tag.Tensor));
+            expect(methodMeta.inputTags).toEqual([et.Tag.Tensor, et.Tag.Tensor, et.Tag.Tensor]);
             module.delete();
         });
 
         test("outputs are tensors", () => {
             const module = et.Module.load("add_mul.pte");
             const methodMeta = module.getMethodMeta("forward");
             expect(methodMeta.outputTags.length).toEqual(1);
-            // Tags can only be checked by strict equality.
-            expect(methodMeta.outputTags[0]).toBe(et.Tag.Tensor);
+            expect(methodMeta.outputTags).toEqual([et.Tag.Tensor]);
             module.delete();
         });
 
@@ -183,8 +178,7 @@ describe("Module", () => {
                 const module = et.Module.load("add_mul.pte");
                 const methodMeta = module.getMethodMeta("forward");
                 methodMeta.inputTensorMeta.forEach((tensorInfo) => {
-                    // ScalarType can only be checked by strict equality.
-                    expect(tensorInfo.scalarType).toBe(et.ScalarType.Float);
+                    expect(tensorInfo.scalarType).toEqual(et.ScalarType.Float);
                 });
                 module.delete();
             });
@@ -311,3 +305,11 @@ describe("Module", () => {
         });
     });
 });
+
+describe("sanity", () => {
+    // Emscripten enums are equal by default for some reason.
+    test("different enums are not equal", () => {
+        expect(et.ScalarType.Float).not.toEqual(et.ScalarType.Long);
+        expect(et.Tag.Int).not.toEqual(et.Tag.Double);
+    });
+});

extension/wasm/wasm_bindings.cpp

@@ -539,6 +539,15 @@ EMSCRIPTEN_BINDINGS(WasmBindings) {
           &JsMethodMeta::memory_planned_buffer_sizes)
       .field("backends", &JsMethodMeta::backends)
       .field("numInstructions", &JsMethodMeta::num_instructions);
+
+// For some reason Embind doesn't make it easy to get the names of enums.
+// Additionally, different enums of the same type are considered to be equal.
+// Assigning the name field fixes both of these issues.
+#define JS_ASSIGN_SCALAR_TYPE_NAME(T, NAME) \
+  EM_ASM(Module.ScalarType.NAME.name = #NAME);
+  JS_FORALL_SUPPORTED_TENSOR_TYPES(JS_ASSIGN_SCALAR_TYPE_NAME)
+#define JS_ASSIGN_TAG_NAME(NAME) EM_ASM(Module.Tag.NAME.name = #NAME);
+  EXECUTORCH_FORALL_TAGS(JS_ASSIGN_TAG_NAME)
 }
 
 } // namespace wasm