Support tuple inputs in extract_submodel

[smpanaro]

In the extract_submodel debugging util, the input reachability check doesn't support ops that take list/tuple inputs e.g. concat, stack

For example, this script:

import coremltools as ct
from coremltools.converters.mil.debugging_utils import extract_submodel
import torch
from torch import nn
import numpy as np

class Net(nn.Module):
    def forward(self, x):
        x = x * x

        chunks = x.chunk(5, dim=-1)
        transformed = []
        for i in range(len(chunks)):
            transformed.append(chunks[i] * i)

        x = torch.cat(transformed, dim=-1)
        x = x ** 0.5
        return x

sample_input = torch.randn(1,32,1,512)
full_model = ct.convert(torch.jit.trace(Net().eval(), sample_input),
                        inputs=[ct.TensorType(shape=sample_input.shape, dtype=np.float16)],
                        minimum_deployment_target=ct.target.iOS16,
                        convert_to="mlprogram")
print("Full model:")
print(full_model._mil_program)

submodel = extract_submodel(full_model, outputs=["var_22"], inputs=["x_cast_fp16"])
print("Submodel:")
print(submodel._mil_program)

On 8.0b1:

Full model:

main[CoreML6](%x_1: (1, 32, 1, 512, fp16)(Tensor)) {
  block0() {
    %x_cast_fp16: (1, 32, 1, 512, fp16)(Tensor) = mul(x=%x_1, y=%x_1, name="x_cast_fp16")
    %var_3_cast_fp16_0: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_1: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_2: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_3: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_4: (1, 32, 1, 100, fp16)(Tensor) = split(x=%x_cast_fp16, split_sizes=[103, 103, 103, 103, 100], axis=-1, name="op_3_cast_fp16")
    %var_9_cast_fp16: (1, 32, 1, 103, fp16)(Tensor) = mul(x=%var_3_cast_fp16_0, y=0.0, name="op_9_cast_fp16")
    %var_13_cast_fp16: (1, 32, 1, 103, fp16)(Tensor) = mul(x=%var_3_cast_fp16_2, y=2.0, name="op_13_cast_fp16")
    %var_15_cast_fp16: (1, 32, 1, 103, fp16)(Tensor) = mul(x=%var_3_cast_fp16_3, y=3.0, name="op_15_cast_fp16")
    %var_17_cast_fp16: (1, 32, 1, 100, fp16)(Tensor) = mul(x=%var_3_cast_fp16_4, y=4.0, name="op_17_cast_fp16")
    %var_20_cast_fp16: (1, 32, 1, 512, fp16)(Tensor) = concat(values=(%var_9_cast_fp16, %var_3_cast_fp16_1, %var_13_cast_fp16, %var_15_cast_fp16, %var_17_cast_fp16), axis=-1, interleave=False, name="op_20_cast_fp16")
    %var_22: (1, 32, 1, 512, fp16)(Tensor) = pow(x=%var_20_cast_fp16, y=0.5, name="op_22_cast_fp16")
  } -> (%var_22)
}

Traceback (most recent call last):
  File "/[removed]/submodel.py", line 28, in <module>
    submodel = extract_submodel(full_model, outputs=["var_22"], inputs=["x_cast_fp16"])
  File "/[removed]/env/lib/python3.10/site-packages/coremltools/converters/mil/debugging_utils.py", line 168, in extract_submodel
    validate_inputs(func, input_vars)
  File "/[removed]/env/lib/python3.10/site-packages/coremltools/converters/mil/debugging_utils.py", line 85, in validate_inputs
    raise ValueError(f"output {output} not reachable from inputs")
ValueError: output var_22 not reachable from inputs

On this branch:

Full model:

main[CoreML6](%x_1: (1, 32, 1, 512, fp16)(Tensor)) {
  block0() {
    %x_cast_fp16: (1, 32, 1, 512, fp16)(Tensor) = mul(x=%x_1, y=%x_1, name="x_cast_fp16")
    %var_3_cast_fp16_0: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_1: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_2: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_3: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_4: (1, 32, 1, 100, fp16)(Tensor) = split(x=%x_cast_fp16, split_sizes=[103, 103, 103, 103, 100], axis=-1, name="op_3_cast_fp16")
    %var_9_cast_fp16: (1, 32, 1, 103, fp16)(Tensor) = mul(x=%var_3_cast_fp16_0, y=0.0, name="op_9_cast_fp16")
    %var_13_cast_fp16: (1, 32, 1, 103, fp16)(Tensor) = mul(x=%var_3_cast_fp16_2, y=2.0, name="op_13_cast_fp16")
    %var_15_cast_fp16: (1, 32, 1, 103, fp16)(Tensor) = mul(x=%var_3_cast_fp16_3, y=3.0, name="op_15_cast_fp16")
    %var_17_cast_fp16: (1, 32, 1, 100, fp16)(Tensor) = mul(x=%var_3_cast_fp16_4, y=4.0, name="op_17_cast_fp16")
    %var_20_cast_fp16: (1, 32, 1, 512, fp16)(Tensor) = concat(values=(%var_9_cast_fp16, %var_3_cast_fp16_1, %var_13_cast_fp16, %var_15_cast_fp16, %var_17_cast_fp16), axis=-1, interleave=False, name="op_20_cast_fp16")
    %var_22: (1, 32, 1, 512, fp16)(Tensor) = pow(x=%var_20_cast_fp16, y=0.5, name="op_22_cast_fp16")
  } -> (%var_22)
}

Running MIL frontend_milinternal pipeline: 0 passes [00:00, ? passes/s]
Running MIL default pipeline: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████| 78/78 [00:00<00:00, 12916.25 passes/s]
Running MIL backend_mlprogram pipeline: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 12/12 [00:00<00:00, 12023.81 passes/s]
/[removed]/env/lib/python3.10/site-packages/coremltools/models/model.py:419: RuntimeWarning: You will not be able to run predict() on this Core ML model. Underlying exception message was: Error compiling model: "compiler error: Error reading protobuf spec. validator error: Description of multiarray feature 'x_cast_fp16' has FLOAT16 dataType, which is only valid in specification version >= 7. This model has version 6".
  _warnings.warn(
Submodel:

main[CoreML6](%x_cast_fp16: (1, 32, 1, 512, fp16)(Tensor)) {
  block0() {
    %var_3_cast_fp16_0: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_1: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_2: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_3: (1, 32, 1, 103, fp16)(Tensor), %var_3_cast_fp16_4: (1, 32, 1, 100, fp16)(Tensor) = split(x=%x_cast_fp16, split_sizes=[103, 103, 103, 103, 100], axis=-1, name="op_3_cast_fp16")
    %var_9_cast_fp16: (1, 32, 1, 103, fp16)(Tensor) = mul(x=%var_3_cast_fp16_0, y=0.0, name="op_9_cast_fp16")
    %var_13_cast_fp16: (1, 32, 1, 103, fp16)(Tensor) = mul(x=%var_3_cast_fp16_2, y=2.0, name="op_13_cast_fp16")
    %var_15_cast_fp16: (1, 32, 1, 103, fp16)(Tensor) = mul(x=%var_3_cast_fp16_3, y=3.0, name="op_15_cast_fp16")
    %var_17_cast_fp16: (1, 32, 1, 100, fp16)(Tensor) = mul(x=%var_3_cast_fp16_4, y=4.0, name="op_17_cast_fp16")
    %var_20_cast_fp16: (1, 32, 1, 512, fp16)(Tensor) = concat(values=(%var_9_cast_fp16, %var_3_cast_fp16_1, %var_13_cast_fp16, %var_15_cast_fp16, %var_17_cast_fp16), axis=-1, interleave=False, name="op_20_cast_fp16")
    %var_22: (1, 32, 1, 512, fp16)(Tensor) = pow(x=%var_20_cast_fp16, y=0.5, name="op_22_cast_fp16")
  } -> (%var_22)
}

[smpanaro]

Not directly related but fixes an error I get with my example conversion script:

RuntimeWarning: You will not be able to run predict() on this Core ML model. Underlying exception message was: Error compiling model: "compiler error: Error reading protobuf spec. validator error: Description of multiarray feature 'x_cast_fp16' has FLOAT16 dataType, which is only valid in specification version >= 7. This model has version 6".

Happy to put this up as a separate PR if that's preferred.

[jakesabathia2]

This is a correct fix :)

[YifanShenSZ]

@jakesabathia2 who added extract_submodel

[jakesabathia2]

Thank you so much @smpanaro for putting this PR!
In order to get this PR merged,
could you also add an unittest for this change?

[smpanaro]

No problem. Just added one (and also fixed one that was failing because of the opset_version change).

[smpanaro]

@jakesabathia2 Let me know if anything else is needed here!