def test_int16_embedding_e2e(
        self, compute_precision, num_embeddings, minimum_deployment_target, symbolic
    ):
        """End-to-end conversion from a torch embedding model."""

        class EmbeddingModel(nn.Module):
            def __init__(self):
                super(EmbeddingModel, self).__init__()
                self.embedding = torch.nn.Embedding(num_embeddings=num_embeddings, embedding_dim=2)

            def forward(self, x):
                return self.embedding(x)

        input_data = np.random.randint(low=0, high=num_embeddings, size=(3, 5))
        input_data = torch.from_numpy(input_data)
        model = EmbeddingModel()
        model.eval()
        traced_model = torch.jit.trace(model, input_data)
        input_shape = (ct.RangeDim(1, 32), ct.RangeDim(1, 32)) if symbolic else input_data.shape
        converted_model = ct.convert(
            traced_model,
            inputs=[ct.TensorType(shape=input_shape, name="input", dtype=np.int32)],
            convert_to="mlprogram",
            compute_precision=compute_precision,
            compute_units=ct.ComputeUnit.CPU_ONLY,
            minimum_deployment_target=minimum_deployment_target,
        )
        prog = converted_model._mil_program

        # The embedding layer is lowered to `gather` op.
        expected_ops = ["gather"]
        if (
            compute_precision == ct.precision.FLOAT16
            and minimum_deployment_target < ct.target.iOS16
        ):
            # Cast from fp16 to fp32 because fp16 is not supported in I/O before iOS16.
            expected_ops.append("cast")
        if (
            minimum_deployment_target >= ct.target.iOS17
            and compute_precision == ct.precision.FLOAT16
            and num_embeddings <= np.iinfo(np.int16).max
        ):
            # The int16 cast only happens for iOS17+ with fp16 precision and there is no overflow.
            expected_ops.insert(0, "cast")
            cast_op = prog["main"].find_ops(op_type="cast")[0]
            assert cast_op.dtype.val == "int16"
            assert cast_op.outputs[0] == prog["main"].find_ops(op_type="gather")[0].indices
        assert get_op_types_in_program(prog) == expected_ops
End-to-end conversion from a torch embedding model.
coremltools/converters/mil/mil/passes/tests/test_quantization_passes.py