[export] ExportedProgram

[angelayi]

class ExportedProgram:
    graph_module: torch.fx.GraphModule
    graph_signature: ExportGraphSignature
    call_spec: CallSpec
    state_dict: Dict[str, Any]    
    symbol_to_range: Dict[sympy.Symbol, Tuple[int, int]]

    @property
    def graph(self):
        return self.graph_module.graph

    def transform(self, *passes: PassType) -> "ExportedProgram":
        # Runs graph based transformations on the given ExportedProgram 
        # and returns a new transformed ExportedProgram
        ...

    def add_runtime_assertions(self) -> "ExportedProgram":
        # Adds runtime assertions based on the constraints
        ...

# Information to maintain user calling/returning specs
@dataclasses.dataclass
class CallSpec:
    in_spec: Optional[pytree.TreeSpec] = None
    out_spec: Optional[pytree.TreeSpec] = None


# Extra information for joint graphs
@dataclasses.dataclass
class ExportBackwardSignature:
    gradients_to_parameters: Dict[str, str]
    gradients_to_user_inputs: Dict[str, str]
    loss_output: str


@dataclasses.dataclass
class ExportGraphSignature:
    parameters: List[str]
    buffers: List[str]

    user_inputs: List[str]
    user_outputs: List[str]
    inputs_to_parameters: Dict[str, str]
    inputs_to_buffers: Dict[str, str]

    buffers_to_mutate: Dict[str, str]

    backward_signature: Optional[ExportBackwardSignature]

Stack from ghstack (oldest at bottom):

• [export] Initial deserialization #102126
• [export] Initial serialization #102125
• [export] Rename graph_module.py to exported_program.py #102260
• -> [export] ExportedProgram #102259

cc @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/102259

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit d57b238:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ydwu4]

Thanks for the rapid pr! This looks good to me.

[avikchaudhuri]

Rebase over #102256
You should be able to grab input_shape_constraints_by_name from https://github.com/pytorch/pytorch/pull/102256/files#diff-cc590ca123eee83ede51a1c290198e331b1e528ea3d616f85cfcfe14afe03122R267

[avikchaudhuri]

I'm not a fan of keeping these private fields. We should be able to eliminate ep._input_shape_constraints in lieu of symbol_to_constraints (generalizing symbol_to_range to include equality constraints, which are basically other symbols).

Also, instead of _input_name_to_example_inputs we should be able to just remember a map from names to the shapes of the example inputs, or even a map from (name, dim) pairs to constant values. Because the only reason we need them is for specialization assertions.

If you still want to keep the example inputs for additional debugging, then sure, keep them...but there should not be any remaining client for it after this change. I believe the right thing to do is to just forget them.

[angelayi]

Chatted with Avik offline -- I'll clean this up in a separate diff.

[avikchaudhuri]

We shouldn't need this function if you move it before export, right?

[avikchaudhuri]

We can move this before export by preprocessing input_name_to_example_inputs

[avikchaudhuri]

Thanks for working on this! This is beginning to shape up to something awesome (no pun intended lol).