Skip to content

Rename a feature inside nested MIL blocks - #2796

Open
LeSingh1 wants to merge 1 commit into
apple:mainfrom
LeSingh1:rename-feature-nested-blocks
Open

Rename a feature inside nested MIL blocks#2796
LeSingh1 wants to merge 1 commit into
apple:mainfrom
LeSingh1:rename-feature-nested-blocks

Conversation

@LeSingh1

@LeSingh1 LeSingh1 commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

rename_feature's ML Program branch walks each function's block_specializations but not the nested blocks that control flow operations hold. cond and while_loop blocks reference outer vars by name, so renaming an input or output that they use leaves those references pointing at the old name. The renamed model then fails to load:

Unable to parse ML Program: at unknown location:
Input 'b' for parameter 'x' does not resolve.

Repro on main:

@mb.program(input_specs=[mb.TensorSpec(shape=(1,), dtype=types.bool),
                         mb.TensorSpec(shape=(1,))])
def prog(a, b):
    def true_fn():
        return mb.add(x=b, y=1.0)
    def false_fn():
        return mb.mul(x=b, y=2.0)
    return mb.cond(pred=mb.squeeze(x=a), _true_fn=true_fn, _false_fn=false_fn)

model = ct.convert(prog, convert_to="mlprogram",
                   minimum_deployment_target=ct.target.iOS16)
spec = model.get_spec()
ct.utils.rename_feature(spec, "b", "renamed_b")
ct.models.MLModel(spec, weights_dir=model.weights_dir).predict(...)  # raises

The block walk is now a recursive helper so nested blocks get the same treatment. Added a test that renames an input used by both cond branches and checks the predictions on each branch; it fails on main with the error above.

rename_feature walked only the top level of each function's block
specializations. Control flow operations such as cond and while_loop hold
nested blocks that reference outer vars by name, so renaming an input or
output of such a model left those references behind and the result no
longer parsed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant