ai-edge-quantizer-nightly 0.0.1.dev20250302__py3-none-any.whl → 0.5.0.dev20260103__py3-none-any.whl

ai_edge_quantizer/transformations/insert_decomposed_hadamard_rotation_test.py

@@ -0,0 +1,244 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Test insertion of the Decomposed Hadamard rotation ops."""
+
+import os
+import numpy as np
+from tensorflow.python.platform import googletest
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import insert_decomposed_hadamard_rotation
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
+
+_TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile('..')
+
+
+class InsertDecomposedHadamardRotationFullyConnectedTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, 'tests/models/single_fc_bias.tflite'
+    )
+    self.model = tfl_flatbuffer_utils.read_model(model_path)
+    self.params = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.ones(1),
+        zero_point=np.zeros(1),
+        hadamard=qtyping.UniformQuantParams.HadamardRotationParams(
+            random_binary_vector=np.ones(1),
+            hadamard_size=2,
+        ),
+    )
+
+  def test_raise_unsupported_qparams(self):
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: 'uniform quantization' in str(err)
+    ):
+      insert_decomposed_hadamard_rotation.insert_decomposed_hadamard_rotation(
+          transformation_utils.TransformationInput(
+              tensor_id=0,
+              op_codes=self.model.operatorCodes,
+              buffers=self.model.buffers,
+              subgraph=self.model.subgraphs[0],
+              producer=-1,
+              consumers=[-1],
+              quant_params=qtyping.NonLinearQuantParams(
+                  num_bits=16, quantized_data=None
+              ),
+          )
+      )
+
+  def test_raise_missing_hadamard_data(self):
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: 'quantization params are not set' in str(err)
+    ):
+      insert_decomposed_hadamard_rotation.insert_decomposed_hadamard_rotation(
+          transformation_utils.TransformationInput(
+              tensor_id=0,
+              op_codes=self.model.operatorCodes,
+              buffers=self.model.buffers,
+              subgraph=self.model.subgraphs[0],
+              producer=-1,
+              consumers=[-1],
+              quant_params=qtyping.UniformQuantParams(
+                  num_bits=8,
+                  quantized_dimension=None,
+                  scale=np.ones(1),
+                  zero_point=np.zeros(1),
+              ),
+          )
+      )
+
+  def test_raise_non_float32_tensor(self):
+    self.model.subgraphs[0].tensors[
+        0
+    ].type = schema_py_generated.TensorType.INT32
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: 'float32 tensors' in str(err)
+    ):
+      insert_decomposed_hadamard_rotation.insert_decomposed_hadamard_rotation(
+          transformation_utils.TransformationInput(
+              tensor_id=0,
+              op_codes=self.model.operatorCodes,
+              buffers=self.model.buffers,
+              subgraph=self.model.subgraphs[0],
+              producer=-1,
+              consumers=[-1],
+              quant_params=self.params,
+          ),
+      )
+
+  def test_insert_decomposed_ops(self):
+    # Insert Decomposed Hadamard ops before fully_connected
+    info = (
+        insert_decomposed_hadamard_rotation.insert_decomposed_hadamard_rotation(
+            transformation_utils.TransformationInput(
+                tensor_id=0,
+                op_codes=self.model.operatorCodes,
+                buffers=self.model.buffers,
+                subgraph=self.model.subgraphs[0],
+                producer=-1,
+                consumers=[0],  # Consumer is the FC op
+                quant_params=self.params,
+            )
+        )
+    )
+    subgraph = self.model.subgraphs[0]
+    self.assertEqual(info.op_id, 0)
+    self.assertEqual(info.num_ops_added, 3)
+    # Model had 4 tensors, added 6 tensors (3 activations 3 constants).
+    self.assertEqual(info.output_tensor_id, 9)
+    self.assertLen(subgraph.tensors, 10)
+    # Model had 1 op code, added RESHAPE and FC.
+    self.assertLen(self.model.operatorCodes, 3)
+    self.assertEqual(
+        self.model.operatorCodes[1].builtinCode,
+        schema_py_generated.BuiltinOperator.RESHAPE,
+    )
+    self.assertEqual(
+        self.model.operatorCodes[2].builtinCode,
+        schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
+    )
+
+    # Op 0: RESHAPE
+    reshape_op = subgraph.operators[0]
+    self.assertEqual(
+        self.model.operatorCodes[reshape_op.opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.RESHAPE,
+    )
+    self.assertEqual(reshape_op.inputs[0], 0)  # Graph input
+    self.assertEqual(reshape_op.outputs[0], 5)  # Reshape output
+
+    # Op 1: FULLY_CONNECTED
+    fc_op = subgraph.operators[1]
+    self.assertEqual(
+        self.model.operatorCodes[fc_op.opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
+    )
+    self.assertEqual(fc_op.inputs[0], 5)  # Reshape output
+    self.assertEqual(fc_op.inputs[1], 6)  # Hadamard matrix tensor
+    self.assertEqual(fc_op.outputs[0], 7)  # FC output
+
+    # Op 2: RESHAPE (post)
+    post_reshape_op = subgraph.operators[2]
+    self.assertEqual(
+        self.model.operatorCodes[post_reshape_op.opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.RESHAPE,
+    )
+    self.assertEqual(post_reshape_op.inputs[0], 7)  # FC output
+    self.assertEqual(post_reshape_op.outputs[0], 9)  # Post Reshape output
+
+    # Op 3: Original FULLY_CONNECTED
+    orig_fc_op = subgraph.operators[3]
+    self.assertEqual(
+        self.model.operatorCodes[orig_fc_op.opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
+    )
+    # Input to the original FC is the post reshape output
+    self.assertEqual(orig_fc_op.inputs[0], 9)
+
+
+class InsertDecomposedHadamardRotationEmbeddingLookupTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, 'tests/models/embedding_lookup.tflite'
+    )
+    self.model = tfl_flatbuffer_utils.read_model(model_path)
+    self.params = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.ones(1),
+        zero_point=np.zeros(1),
+        hadamard=qtyping.UniformQuantParams.HadamardRotationParams(
+            random_binary_vector=np.ones(1),
+            hadamard_size=2,
+        ),
+    )
+
+  def test_insert_decomposed_ops(self):
+    # Insert Decomposed Hadamard ops after embedding_lookup
+    info = (
+        insert_decomposed_hadamard_rotation.insert_decomposed_hadamard_rotation(
+            transformation_utils.TransformationInput(
+                tensor_id=2,  # Output of embedding_lookup
+                op_codes=self.model.operatorCodes,
+                buffers=self.model.buffers,
+                subgraph=self.model.subgraphs[0],
+                producer=0,
+                consumers=[-1],  # Output is a graph output
+                quant_params=self.params,
+            )
+        )
+    )
+    subgraph = self.model.subgraphs[0]
+    self.assertEqual(info.op_id, 1)
+    self.assertEqual(info.num_ops_added, 3)
+    # Model had 3 tensors, added 6 (3 activations 3 constants).
+    self.assertEqual(info.output_tensor_id, 8)
+    self.assertLen(subgraph.tensors, 9)
+    # Model had 1 op code, added RESHAPE and FC.
+    self.assertLen(self.model.operatorCodes, 3)
+
+    # Op 0: EMBEDDING_LOOKUP (Original)
+    # Op 1: RESHAPE
+    reshape_op = subgraph.operators[1]
+    self.assertEqual(reshape_op.inputs[0], 2)  # Embedding lookup output
+    self.assertEqual(reshape_op.outputs[0], 4)
+
+    # Op 2: FULLY_CONNECTED
+    fc_op = subgraph.operators[2]
+    self.assertEqual(fc_op.inputs[0], 4)
+    self.assertEqual(fc_op.inputs[1], 5)  # Hadamard matrix
+    self.assertEqual(fc_op.outputs[0], 6)
+
+    # Op 3: RESHAPE (post)
+    post_reshape_op = subgraph.operators[3]
+    self.assertEqual(post_reshape_op.inputs[0], 6)
+    self.assertEqual(post_reshape_op.outputs[0], 8)
+
+    # Check graph output
+    self.assertIn(8, subgraph.outputs)
+    self.assertNotIn(2, subgraph.outputs)
+
+
+if __name__ == '__main__':
+  googletest.main()

ai_edge_quantizer/transformations/insert_hadamard_rotation.py

@@ -0,0 +1,186 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Hadamard rotation pattern transformation."""
+
+from flatbuffers import flexbuffers
+import numpy as np
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
+
+
+def _to_flexbuffer(
+    hadamard_size: int,
+    random_binary_vector: list[np.int8],
+) -> bytes:
+  """Converts hadamard_size to flexbuffer."""
+  fbb = flexbuffers.Builder()
+  with fbb.Map():
+    fbb.Int('hadamard_size', hadamard_size)
+    fbb.VectorFromElements('random_binary_vector', random_binary_vector)
+  return fbb.Finish()
+
+
+def _update_embedding_lookup_consumers(
+    transformation: transformation_utils.TransformationInput,
+    new_tensor_id: int,
+) -> bool:
+  """Updates the consumers of the embedding lookup op to use the new tensor.
+
+  Args:
+    transformation: The transformation input to update the consumers of.
+    new_tensor_id: The new tensor id to use as the input to the embedding lookup
+      consumers.
+  """
+  for consumer in transformation.consumers:
+    # If the consumer is a graph output and not an op, we can ignore it here
+    # since the graph output will be updated later.
+    if consumer == -1:
+      continue
+    consumer_op = transformation.subgraph.operators[consumer]
+    # Find the input that was attached to the insertion point, and replace it
+    # with the new tensor.
+    for i in range(len(consumer_op.inputs)):
+      if consumer_op.inputs[i] == transformation.tensor_id:
+        consumer_op.inputs[i] = new_tensor_id
+
+
+def _update_fully_connected_consumers(
+    transformation: transformation_utils.TransformationInput,
+    new_tensor_id: int,
+) -> bool:
+  """Updates the fully connected op(s) to use the new tensor.
+
+  Since the new tensor is inserted to the fully_connected's input, we need to
+  scan each consumer (in case of multiple fully_connected ops), and update
+  the input tensor to the new tensor.
+
+  Args:
+    transformation: The transformation input to update the consumers of.
+    new_tensor_id: The new tensor id to use as the input to the fully connected
+      consumers.
+
+  Returns:
+    True if the fully connected op(s) were updated to use the new tensor.
+  """
+  updated = False
+  for consumer in transformation.consumers:
+    if (
+        transformation_utils.get_schema_op_id(transformation, consumer)
+        == schema_py_generated.BuiltinOperator.FULLY_CONNECTED
+    ):
+      transformation.subgraph.operators[consumer].inputs[0] = new_tensor_id
+      updated = True
+  return updated
+
+
+def insert_hadamard_rotation(
+    transformation_input: transformation_utils.TransformationInput,
+) -> qtyping.TransformationInfo:
+  """Inserts a custom aeq.hadamard_rotation op on this tensor.
+
+  This function works for float32 tensors only.
+
+  Args:
+    transformation_input: The transformation input to insert the custom op on.
+
+  Returns:
+    The transformation info of the inserted custom op.
+
+  Raises:
+    ValueError: If the transformation input is not a uniform quantization
+    transformation.
+    ValueError: If the Hadamard quantization params are not set.
+    ValueError: If the tensor is not a float32 tensor.
+    ValueError: If no supported ops were found as the tensor's producer or
+    consumers.
+  """
+  if not isinstance(
+      transformation_input.quant_params, qtyping.UniformQuantParams
+  ):
+    raise ValueError('Hadamard rotation supports uniform quantization only')
+
+  if transformation_input.quant_params.hadamard is None:
+    raise ValueError(
+        'Hadamard rotation quantization params are not set but op insertion is'
+        ' requested.'
+    )
+
+  tensor = transformation_input.subgraph.tensors[transformation_input.tensor_id]
+  if tensor.type != schema_py_generated.TensorType.FLOAT32:
+    raise ValueError(
+        'The Hadamard rotation op supports float32 tensors only. Got'
+        f' {tensor.type} tensor.'
+    )
+
+  # Create new custom op with the current tensor as input and a new activation
+  # tensor as output.
+  custom_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.CUSTOM,
+      transformation_input.op_codes,
+      'aeq.hadamard_rotation',
+  )
+  custom_op = schema_py_generated.OperatorT()
+  custom_op.opcodeIndex = custom_op_code_idx
+  custom_op.inputs = [transformation_input.tensor_id]
+  custom_op.customOptions = _to_flexbuffer(
+      transformation_input.quant_params.hadamard.hadamard_size,
+      transformation_input.quant_params.hadamard.random_binary_vector.tolist(),
+  )
+  new_tensor_id = transformation_utils.add_new_activation_tensor(
+      tensor.name + b'_rotated',
+      tensor.shapeSignature
+      if tensor.shapeSignature is not None
+      else tensor.shape,
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+  )
+  custom_op.outputs = [new_tensor_id]
+
+  # Update the users of this tensor to use the new tensor.
+  if (
+      transformation_utils.get_producer_schema_op_id(transformation_input)
+      == schema_py_generated.BuiltinOperator.EMBEDDING_LOOKUP
+  ):
+    _update_embedding_lookup_consumers(transformation_input, new_tensor_id)
+  elif not _update_fully_connected_consumers(
+      transformation_input, new_tensor_id
+  ):
+    raise ValueError(
+        'The Hadamard rotation op supports embedding lookup and fully connected'
+        ' ops only, but no such ops were found.'
+    )
+
+  # If the tensor is a graph output, we need to replace the tensor with the
+  # new tensor.
+  for i, output in enumerate(transformation_input.subgraph.outputs):
+    if output == transformation_input.tensor_id:
+      transformation_input.subgraph.outputs[i] = new_tensor_id
+
+  # Find the actual insertion point. The insertion point should be after the
+  # producer op and before the first consumer op. The max() operation ensures
+  # that we're not using -1 as the insertion point.
+  first_consumer_id = min(transformation_input.consumers)
+  op_id = max(transformation_input.producer + 1, first_consumer_id)
+
+  # Insert the custom op.
+  transformation_input.subgraph.operators.insert(op_id, custom_op)
+
+  return qtyping.TransformationInfo(
+      op_id=op_id,
+      num_ops_added=1,
+      output_tensor_id=new_tensor_id,
+  )

ai_edge_quantizer/transformations/insert_hadamard_rotation_test.py

@@ -0,0 +1,200 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Test insertion of the Hadamard rotation custom op."""
+
+import os
+import numpy as np
+from tensorflow.python.platform import googletest
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import insert_hadamard_rotation
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
+
+_TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile('..')
+
+
+class InsertHadamardRotationFullyConnectedTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, 'tests/models/single_fc_bias.tflite'
+    )
+    self.model = tfl_flatbuffer_utils.read_model(model_path)
+    self.params = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.ones(1),
+        zero_point=np.zeros(1),
+        hadamard=qtyping.UniformQuantParams.HadamardRotationParams(
+            random_binary_vector=np.ones(1),
+            hadamard_size=2,
+        ),
+    )
+
+  def test_raise_unsupported_qparams(self):
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: 'uniform quantization' in str(err)
+    ):
+      insert_hadamard_rotation.insert_hadamard_rotation(
+          transformation_utils.TransformationInput(
+              tensor_id=0,
+              op_codes=self.model.operatorCodes,
+              buffers=self.model.buffers,
+              subgraph=self.model.subgraphs[0],
+              producer=-1,
+              consumers=[-1],
+              quant_params=qtyping.NonLinearQuantParams(
+                  num_bits=16, quantized_data=None
+              ),
+          )
+      )
+
+  def test_raise_missing_hadamard_data(self):
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: 'quantization params are not set' in str(err)
+    ):
+      insert_hadamard_rotation.insert_hadamard_rotation(
+          transformation_utils.TransformationInput(
+              tensor_id=0,
+              op_codes=self.model.operatorCodes,
+              buffers=self.model.buffers,
+              subgraph=self.model.subgraphs[0],
+              producer=-1,
+              consumers=[-1],
+              quant_params=qtyping.UniformQuantParams(
+                  num_bits=8,
+                  quantized_dimension=None,
+                  scale=np.ones(1),
+                  zero_point=np.zeros(1),
+              ),
+          )
+      )
+
+  def test_raise_non_float32_tensor(self):
+    self.model.subgraphs[0].tensors[
+        0
+    ].type = schema_py_generated.TensorType.INT32
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: 'float32 tensors' in str(err)
+    ):
+      insert_hadamard_rotation.insert_hadamard_rotation(
+          transformation_utils.TransformationInput(
+              tensor_id=0,
+              op_codes=self.model.operatorCodes,
+              buffers=self.model.buffers,
+              subgraph=self.model.subgraphs[0],
+              producer=-1,
+              consumers=[-1],
+              quant_params=self.params,
+          ),
+      )
+
+  def test_insert_single_custom_op(self):
+    # Insert aeq.hadamard_rotation before fully_connected
+    info = insert_hadamard_rotation.insert_hadamard_rotation(
+        transformation_utils.TransformationInput(
+            tensor_id=0,
+            op_codes=self.model.operatorCodes,
+            buffers=self.model.buffers,
+            subgraph=self.model.subgraphs[0],
+            producer=-1,
+            consumers=[-1],
+            quant_params=self.params,
+        )
+    )
+    subgraph = self.model.subgraphs[0]
+    self.assertEqual(info.op_id, 0)
+    self.assertEqual(info.num_ops_added, 1)
+    # Model had 4 tensors, added 1.
+    self.assertEqual(info.output_tensor_id, 4)
+    self.assertLen(subgraph.tensors, 5)
+    # Model had 1 op, added a new one.
+    self.assertLen(self.model.operatorCodes, 2)
+    self.assertEqual(
+        self.model.operatorCodes[1].builtinCode,
+        schema_py_generated.BuiltinOperator.CUSTOM,
+    )
+    # First op is now the custom op, precedes fully_connected.
+    self.assertEqual(
+        self.model.operatorCodes[subgraph.operators[0].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.CUSTOM,
+    )
+    # Input to the custom op is graph input
+    self.assertEqual(subgraph.operators[0].inputs[0], 0)
+    # Input to the FC is the custom op output
+    self.assertEqual(subgraph.operators[1].inputs[0], 4)
+
+
+class InsertHadamardRotationEmbeddingLookupTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, 'tests/models/embedding_lookup.tflite'
+    )
+    self.model = tfl_flatbuffer_utils.read_model(model_path)
+    self.params = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.ones(1),
+        zero_point=np.zeros(1),
+        hadamard=qtyping.UniformQuantParams.HadamardRotationParams(
+            random_binary_vector=np.ones(1),
+            hadamard_size=2,
+        ),
+    )
+
+  def test_insert_single_custom_op(self):
+    # Insert aeq.hadamard_rotation after embedding_lookup
+    info = insert_hadamard_rotation.insert_hadamard_rotation(
+        transformation_utils.TransformationInput(
+            tensor_id=2,
+            op_codes=self.model.operatorCodes,
+            buffers=self.model.buffers,
+            subgraph=self.model.subgraphs[0],
+            producer=0,
+            consumers=[-1],
+            quant_params=self.params,
+        )
+    )
+    subgraph = self.model.subgraphs[0]
+    self.assertEqual(info.op_id, 1)
+    self.assertEqual(info.num_ops_added, 1)
+    # Model had 3 tensors, added 1.
+    self.assertEqual(info.output_tensor_id, 3)
+    self.assertLen(subgraph.tensors, 4)
+    # Model had 1 op, added a new one.
+    self.assertLen(self.model.operatorCodes, 2)
+    self.assertEqual(
+        self.model.operatorCodes[1].builtinCode,
+        schema_py_generated.BuiltinOperator.CUSTOM,
+    )
+    # Second op is now the custom op, after embedding_lookup.
+    self.assertEqual(
+        self.model.operatorCodes[subgraph.operators[1].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.CUSTOM,
+    )
+    # Input to the custom op is embedding's output
+    self.assertEqual(subgraph.operators[1].inputs[0], 2)
+    # Custom op's output is the new tensor
+    self.assertEqual(subgraph.operators[1].outputs[0], 3)
+
+
+if __name__ == '__main__':
+  googletest.main()