Enable explicit input/output naming and shape specs in @mb.program decorator
- Dominant language
- Python
- Stars
- 5.4k
- Forks
- 850
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 10
Description
## 🌱 Describe your Feature Request
- Extend the @mb.program decorator to explicitly process name parameters for input_specs on each TensorSpec and StateTensorSpec, and introduce an output_specs argument so that output feature names and shapes can be declared at graph‐construction time. This would allow users to control MLModel feature names and preserve FlexibleShapeInformation for both inputs and outputs without post‐conversion or inline python naming workarounds.
## How can this feature be used?
- Defining a transformer MIL graph where you need to name the 1×L input_ids tensor, the KV‐cache states (k_state, v_state, offset), and the two outputs (logits_tok, done) directly in the decorator.
- Generating consistent Swift/Obj-C bindings with predictable feature names, simplifying the autoregressive inference loop in client code.
## Describe alternatives you've considered
- Passing inputs=[…] and outputs=[…] to ct.convert(...) after building the MIL graph—this only affects conversions from source models e.g. from pyTorch, not hand‐crafted MIL graphs.
- Accepting the current default naming and manually renaming features in downstream code.
## Additional context
- We hand‐construct the MIL graph with @mb.program with passing input_specs and opset_version where input_specs contains elements of type TensorSpec and StateTensorSpec, then call ct.convert(graph, convert_to="mlprogram", …) without any effect on the names of the inputs_specs.
- A decorator‐based specifier for input/output names and shapes would streamline automated pipelines and remove friction in dynamic-model workflows.
Contributor guide
Research direction
Start at the @mb.program decorator and its handling of TensorSpec and StateTensorSpec input_specs, then trace how ct.convert processes hand-constructed MIL graphs. Define the expected output_specs behavior and verify that declared input and output names and shapes, including flexible shape information, are preserved in the resulting MLModel features.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100