Update embedding_lookup model with signature

PiperOrigin-RevId: 755525455

Author: paulinesho

ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting_test.py

@@ -531,9 +531,9 @@ def test_embedding_lookup_weight_only_succeeds(self):
 
     op_tensor_names = {}
     op_tensor_names["weight"] = (
-        "jax2tf_export_func_/...y_yz-_...z/pjit__einsum_/MatMul;jax2tf_export_func_/pjit__one_hot_/Equal;jax2tf_export_func_/pjit__one_hot_/Cast_1"
+        "jit(export_func)/jit(main)/...y,yz->...z/dot_general;jit(export_func)/jit(main)/jit(_one_hot)/eq;jit(export_func)/jit(main)/jit(_one_hot)/convert_element_type"
     )
-    op_tensor_names["input"] = "inputs"
+    op_tensor_names["input"] = "lookup"
     op_tensor_names["output"] = "Identity_1"
 
     # TODO: b/335913710 - Rename the test function.

ai_edge_quantizer/algorithms/uniform_quantize/op_architecture_tests/embedding_lookup_test.py

@@ -93,9 +93,9 @@ def test_embedding_lookup_succeeds(
     op = subgraph0.operators[subgraph_op_id]
     op_tensor_names = {}
     op_tensor_names["weight"] = (
-        "jax2tf_export_func_/...y_yz-_...z/pjit__einsum_/MatMul;jax2tf_export_func_/pjit__one_hot_/Equal;jax2tf_export_func_/pjit__one_hot_/Cast_1"
+        "jit(export_func)/jit(main)/...y,yz->...z/dot_general;jit(export_func)/jit(main)/jit(_one_hot)/eq;jit(export_func)/jit(main)/jit(_one_hot)/convert_element_type"
     )
-    op_tensor_names["input"] = "inputs"
+    op_tensor_names["input"] = "lookup"
     op_tensor_names["output"] = "Identity_1"
     self._op_test_info.op_tensor_names = op_tensor_names
     self._op_test_info.quantized_dimension = 0

ai_edge_quantizer/tests/end_to_end_tests/embedding_lookup_test.py

@@ -52,18 +52,20 @@ def setUp(self):
     self._quantizer = quantizer.Quantizer(self.float_model_path)
 
   @parameterized.parameters(
-      '../../recipes/default_af32w8float_recipe.json',
-      '../../recipes/default_af32w4float_recipe.json',
-      '../../recipes/dynamic_legacy_wi8_afp32_recipe.json',
-      '../../recipes/dynamic_wi8_afp32_recipe.json',
+      ('../../recipes/default_af32w8float_recipe.json', 1700),
+      ('../../recipes/default_af32w4float_recipe.json', 1600),
+      ('../../recipes/dynamic_legacy_wi8_afp32_recipe.json', 1400),
+      ('../../recipes/dynamic_wi8_afp32_recipe.json', 1400),
   )
-  def test_embedding_lookup_model_int_weight_only(self, recipe_path):
+  def test_embedding_lookup_model_int_weight_only(
+      self, recipe_path, expected_model_size
+  ):
     recipe_path = test_utils.get_path_to_datafile(recipe_path)
     self._quantizer.load_quantization_recipe(recipe_path)
     self.assertFalse(self._quantizer.need_calibration)
     quant_result = self._quantizer.quantize()
     # Check model size.
-    self.assertLess(len(quant_result.quantized_model), 2000)
+    self.assertLess(len(quant_result.quantized_model), expected_model_size)
 
     # TODO: b/364405203 - Enable after 0 signature works.
     # comparison_result = self._quantizer.validate(
@@ -91,8 +93,7 @@ def test_embedding_lookup_model_fp16_weight_only(self):
         ),
     )
     quant_result = self._quantizer.quantize()
-    print(len(quant_result.quantized_model))
-    self.assertLess(len(quant_result.quantized_model), 2000)
+    self.assertLess(len(quant_result.quantized_model), 1600)
 
     # TODO: b/364405203 - Enable after 0 signature works.
     # comparion_result = self._quantizer.validate(
@@ -106,18 +107,20 @@ def test_embedding_lookup_model_fp16_weight_only(self):
     # )
 
   @parameterized.parameters(
-      '../../recipes/default_a8w8_recipe.json',
-      '../../recipes/default_a16w8_recipe.json',
+      ('../../recipes/default_a8w8_recipe.json', 1400),
+      ('../../recipes/default_a16w8_recipe.json', 1400),
   )
-  def test_embedding_lookup_model_full_integer(self, recipe_path):
+  def test_embedding_lookup_model_full_integer(
+      self, recipe_path, expected_model_size
+  ):
     calibration_result = {
         'Identity_1': {'min': -2.0, 'max': 2.0},
     }
     recipe_path = test_utils.get_path_to_datafile(recipe_path)
     self._quantizer.load_quantization_recipe(recipe_path)
     self.assertTrue(self._quantizer.need_calibration)
     quant_result = self._quantizer.quantize(calibration_result)
-    self.assertLess(len(quant_result.quantized_model), 2000)
+    self.assertLess(len(quant_result.quantized_model), expected_model_size)
     # TODO: b/364405203 - Enable after 0 signature works.
     # comparion_result = self._quantizer.validate(
     #     error_metrics='mse',